Update main.py

main.py

@@ -36,29 +36,33 @@ def process_image(path):
     out = []
     for r in results:
         for box in r.boxes:
-            cls = int(box.cls)
-            conf = float(box.conf)
+            cls   = int(box.cls)
+            conf  = float(box.conf)
             xywhn = box.xywhn[0].tolist()
             lat, lng = random_riyadh()
             out.append({
                 "damage_type": CLASS_NAMES.get(cls, "other"),
-                "confidence": round(conf, 3),
-                "severity": severity(conf, xywhn[2] * xywhn[3]),
-                "bbox": xywhn,
-                "frame": 0,
-                "latitude": lat,
-                "longitude": lng,
+                "confidence" : round(conf, 3),
+                "severity"   : severity(conf, xywhn[2] * xywhn[3]),
+                "bbox"       : xywhn,
+                "frame"      : 0,
+                "latitude"   : lat,
+                "longitude"  : lng,
             })
     return out
 
 
 def process_video(path):
     results = model.track(
-    source=path, conf=0.25, tracker="bytetrack.yaml",
-    stream=True, verbose=False, save=True,
-)
+        source=path,
+        conf=0.25,
+        tracker="bytetrack.yaml",
+        stream=True,
+        verbose=False,
+        save=False,
+    )
 
-    seen = {}  # track_id -> best detection
+    seen = {}
 
     for frame_idx, r in enumerate(results):
         if r.boxes is None or r.boxes.id is None:
@@ -71,15 +75,16 @@ def process_video(path):
             r.boxes.xywhn.tolist(),
         ):
             if tid not in seen or conf > seen[tid]["confidence"]:
-                lat, lng = random_riyadh() if tid not in seen else (seen[tid]["latitude"], seen[tid]["longitude"])
+                lat = seen[tid]["latitude"] if tid in seen else random_riyadh()[0]
+                lng = seen[tid]["longitude"] if tid in seen else random_riyadh()[1]
                 seen[tid] = {
                     "damage_type": CLASS_NAMES.get(cls, "other"),
-                    "confidence": round(conf, 3),
-                    "severity": severity(conf, xywhn[2] * xywhn[3]),
-                    "bbox": xywhn,
-                    "frame": frame_idx,
-                    "latitude": lat,
-                    "longitude": lng,
+                    "confidence" : round(conf, 3),
+                    "severity"   : severity(conf, xywhn[2] * xywhn[3]),
+                    "bbox"       : xywhn,
+                    "frame"      : frame_idx,
+                    "latitude"   : lat,
+                    "longitude"  : lng,
                 }
 
     return list(seen.values())
@@ -93,7 +98,7 @@ def root():
 @app.post("/detect")
 async def detect(file: UploadFile = File(...)):
     contents = await file.read()
-    suffix = "." + file.filename.split(".")[-1]
+    suffix   = "." + file.filename.split(".")[-1]
 
     with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
         tmp.write(contents)
@@ -104,4 +109,4 @@ async def detect(file: UploadFile = File(...)):
     finally:
         os.unlink(tmp_path)
 
-    return {"total": len(detections), "detections": detections}
+    return {"total": len(detections), "detections": detections}