Update app.py

app.py

@@ -1,57 +1,57 @@
-from fastapi import FastAPI, File, UploadFile
-from pydantic import BaseModel
-import uvicorn
-import cv2
-import numpy as np
-import torch
-from detectron2.engine import DefaultPredictor
-from detectron2.config import get_cfg
-from detectron2 import model_zoo
-from detectron2.data import MetadataCatalog
-
-# Create FastAPI app
-app = FastAPI()
-
-# Set up the Detectron2 model
-cfg = get_cfg()
-cfg.merge_from_file(model_zoo.get_config_file("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"))
-cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml")
-cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set a threshold for detection
-predictor = DefaultPredictor(cfg)
-
-class PredictionResponse(BaseModel):
-    objects: dict
-
-# API endpoint to process the image and return pixel coordinates
-@app.post("/predict/", response_model=PredictionResponse)
-async def predict(file: UploadFile = File(...)):
-    contents = await file.read()
-
-    # Load the image from bytes
-    nparr = np.frombuffer(contents, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-
-    # Make prediction
-    outputs = predictor(img)
-
-    # Process outputs to get pixel coordinates for each object
-    panoptic_seg, segments_info = outputs["panoptic_seg"]
-
-    objects_coordinates = {}
-    
-    # Iterate over each object and collect coordinates
-    for segment in segments_info:
-        category_id = segment["category_id"]
-        mask = panoptic_seg == segment["id"]
-        coordinates = np.argwhere(mask.cpu().numpy())
-
-        # Convert category_id into a human-readable label
-        label = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes[category_id]
-
-        objects_coordinates[label] = coordinates.tolist()  # Convert to a list of coordinates
-
-    return {"objects": objects_coordinates}
-
-# Start the API
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+from fastapi import FastAPI, File, UploadFile
+from pydantic import BaseModel
+import uvicorn
+import cv2
+import numpy as np
+import torch
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2 import model_zoo
+from detectron2.data import MetadataCatalog
+
+# Create FastAPI app
+app = FastAPI()
+
+# Set up the Detectron2 model
+cfg = get_cfg()
+cfg.merge_from_file(model_zoo.get_config_file("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"))
+cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml")
+cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set a threshold for detection
+predictor = DefaultPredictor(cfg)
+
+class PredictionResponse(BaseModel):
+    objects: dict
+
+# API endpoint to process the image and return pixel coordinates
+@app.post("/predict/", response_model=PredictionResponse)
+async def predict(file: UploadFile = File(...)):
+    contents = await file.read()
+
+    # Load the image from bytes
+    nparr = np.frombuffer(contents, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+    # Make prediction
+    outputs = predictor(img)
+
+    # Process outputs to get pixel coordinates for each object
+    panoptic_seg, segments_info = outputs["panoptic_seg"]
+
+    objects_coordinates = {}
+    
+    # Iterate over each object and collect coordinates
+    for segment in segments_info:
+        category_id = segment["category_id"]
+        mask = panoptic_seg == segment["id"]
+        coordinates = np.argwhere(mask.cpu().numpy())
+
+        # Convert category_id into a human-readable label
+        label = MetadataCatalog.get(cfg.DATASETS.TRAIN[0]).thing_classes[category_id]
+
+        objects_coordinates[label] = coordinates.tolist()  # Convert to a list of coordinates
+
+    return {"objects": objects_coordinates}
+
+# Start the API
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)