DurgaDeepak commited on
Commit
595fad2
·
verified ·
1 Parent(s): 86f4d8f

Update core/process.py

Browse files
Files changed (1) hide show
  1. core/process.py +35 -0
core/process.py CHANGED
@@ -1,6 +1,41 @@
1
  import logging
2
  import time
3
  import timeout_decorator
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4
 
5
  @timeout_decorator.timeout(35, use_signals=False) # 35 sec limit per image
6
  def process_image(
 
1
  import logging
2
  import time
3
  import timeout_decorator
4
+ import io
5
+ import zipfile
6
+ import json
7
+ import cv2
8
+ from PIL import Image
9
+
10
+ from registry import get_model
11
+ from core.describe_scene import describe_scene
12
+ from utils.helpers import generate_session_id, log_runtime
13
+ # Setup logging
14
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
15
+ logger = logging.getLogger(__name__)
16
+
17
+ # Model mappings
18
+ DETECTION_MODEL_MAP = {
19
+ "YOLOv5-Nano": "yolov5n-seg",
20
+ "YOLOv5-Small": "yolov5s-seg",
21
+ "YOLOv8-Small": "yolov8s",
22
+ "YOLOv8-Large": "yolov8l",
23
+ "RT-DETR": "rtdetr" # For future support
24
+ }
25
+
26
+ SEGMENTATION_MODEL_MAP = {
27
+ "SegFormer-B0": "nvidia/segformer-b0-finetuned-ade-512-512",
28
+ "SegFormer-B5": "nvidia/segformer-b5-finetuned-ade-512-512",
29
+ "DeepLabV3-ResNet50": "deeplabv3_resnet50"
30
+ }
31
+
32
+ DEPTH_MODEL_MAP = {
33
+ "MiDaS v21 Small 256": "midas_v21_small_256",
34
+ "MiDaS v21 384": "midas_v21_384",
35
+ "DPT Hybrid 384": "dpt_hybrid_384",
36
+ "DPT Swin2 Large 384": "dpt_swin2_large_384",
37
+ "DPT Beit Large 512": "dpt_beit_large_512"
38
+ }
39
 
40
  @timeout_decorator.timeout(35, use_signals=False) # 35 sec limit per image
41
  def process_image(