DurgaDeepak commited on
Commit
eb9703e
·
verified ·
1 Parent(s): b9a2f92

Update core/process.py

Browse files
Files changed (1) hide show
  1. core/process.py +7 -5
core/process.py CHANGED
@@ -18,16 +18,17 @@ logger = logging.getLogger(__name__)
18
 
19
  # Model mappings
20
  DETECTION_MODEL_MAP = {
21
- "YOLOv5-Nano": "yolov5n-seg",
22
- "YOLOv5-Small": "yolov5s-seg",
23
  "YOLOv8-Small": "yolov8s",
24
  "YOLOv8-Large": "yolov8l",
25
- "RT-DETR": "rtdetr" # For future support
 
26
  }
27
 
28
  SEGMENTATION_MODEL_MAP = {
29
- "SegFormer-B0": "nvidia/segformer-b0-finetuned-ade-512-512",
30
- "SegFormer-B5": "nvidia/segformer-b5-finetuned-ade-512-512",
31
  "DeepLabV3-ResNet50": "deeplabv3_resnet50"
32
  }
33
 
@@ -39,6 +40,7 @@ DEPTH_MODEL_MAP = {
39
  "DPT Beit Large 512": "dpt_beit_large_512"
40
  }
41
 
 
42
  @timeout_decorator.timeout(35, use_signals=False) # 35 sec limit per image
43
  def process_image(
44
  image: Image.Image,
 
18
 
19
  # Model mappings
20
  DETECTION_MODEL_MAP = {
21
+ "YOLOv5-Nano": "yolov5n",
22
+ "YOLOv5-Small": "yolov5s",
23
  "YOLOv8-Small": "yolov8s",
24
  "YOLOv8-Large": "yolov8l",
25
+ "YOLOv11-Beta": "yolov11b",
26
+ "RT-DETR": "rtdetr"
27
  }
28
 
29
  SEGMENTATION_MODEL_MAP = {
30
+ "SegFormer-B0": "segformer_b0",
31
+ "SegFormer-B5": "segformer_b5",
32
  "DeepLabV3-ResNet50": "deeplabv3_resnet50"
33
  }
34
 
 
40
  "DPT Beit Large 512": "dpt_beit_large_512"
41
  }
42
 
43
+
44
  @timeout_decorator.timeout(35, use_signals=False) # 35 sec limit per image
45
  def process_image(
46
  image: Image.Image,