Update predict.py

predict.py

@@ -10,6 +10,9 @@ from scipy import ndimage
 from ultralytics import YOLO
 from PIL import Image
 import io
+import threading
+
+live_view_running = False
 
 app = FastAPI()
 uploads_dir = 'uploads'
@@ -56,13 +59,33 @@ def calculate_wound_dimensions(mask):
     area_cm2 = length_cm * breadth_cm
     return length_cm, breadth_cm, depth_cm, area_cm2
 
+# Draw YOLO detection landmarks (bounding boxes) on the image
+def draw_square_landmarks(frame):
+    results = yolo_model(frame)[0]
+    for box in results.boxes.xyxy.tolist():
+        x1, y1, x2, y2 = map(int, box)
+        w = x2 - x1
+        h = y2 - y1
+        side = max(w, h)
+        cx = x1 + w // 2
+        cy = y1 + h // 2
+        new_x1 = max(cx - side // 2, 0)
+        new_y1 = max(cy - side // 2, 0)
+        new_x2 = new_x1 + side
+        new_y2 = new_y1 + side
+        cv2.rectangle(frame, (new_x1, new_y1), (new_x2, new_y2), (0, 255, 0), 2)
+    return frame
+
 @app.post("/analyze_wound")
 async def analyze_wounds(file: UploadFile = File(...)):
     if file.filename.lower().endswith(('.png', '.jpg', '.jpeg')):
         contents = await file.read()
-        nparr = np.fromstring(contents, np.uint8)
+        nparr = np.frombuffer(contents, np.uint8)  # safer than np.fromstring
         img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
         results = yolo_model(img) 
+        img = draw_square_landmarks(img)  # ✅ Add landmarks to the original image
+
         combined_xmin = float('inf')
         combined_ymin = float('inf')
         combined_xmax = float('-inf')
@@ -73,29 +96,37 @@ async def analyze_wounds(file: UploadFile = File(...)):
             combined_ymin = min(combined_ymin, ymin)
             combined_xmax = max(combined_xmax, xmax)
             combined_ymax = max(combined_ymax, ymax)
+
         combined_xmin = int(combined_xmin)
         combined_ymin = int(combined_ymin)
         combined_xmax = int(combined_xmax)
         combined_ymax = int(combined_ymax)
+
         combined_img = img[combined_ymin:combined_ymax, combined_xmin:combined_xmax]
         combined_img_resized = cv2.resize(combined_img, (224, 224))
         img_array = img_to_array(combined_img_resized) / 255.0
         img_array = np.expand_dims(img_array, axis=0)
+
         output = segmentation_model.predict(img_array)
         predicted_mask = output[0]
+
         mask_overlay = (predicted_mask.squeeze() * 255).astype(np.uint8)
         mask_overlay_colored = np.zeros((mask_overlay.shape[0], mask_overlay.shape[1], 3), dtype=np.uint8)
         mask_overlay_colored[mask_overlay > 200] = [255, 0, 0]  # Red
         mask_overlay_colored[(mask_overlay > 100) & (mask_overlay <= 200)] = [0, 255, 0]  # Green
         mask_overlay_colored[mask_overlay <= 100] = [0, 0, 255]  # Blue
+
         mask_overlay_colored = cv2.resize(mask_overlay_colored, (224, 224))
         blended_image = cv2.addWeighted(combined_img_resized.astype(np.uint8), 0.6, mask_overlay_colored, 0.4, 0)
+
         segmented_image = Image.fromarray(cv2.cvtColor(blended_image, cv2.COLOR_BGR2RGB))
         img_byte_arr = io.BytesIO()
         segmented_image.save(img_byte_arr, format='PNG')
         img_byte_arr.seek(0)
+
         length_cm, breadth_cm, depth_cm, area_cm2 = calculate_wound_dimensions(predicted_mask)
         moisture = calculate_moisture_and_texture(combined_img)
+
         response = Response(img_byte_arr.getvalue(), media_type='image/png')
         response.headers['X-Length-Cm'] = str(length_cm)
         response.headers['X-Breadth-Cm'] = str(breadth_cm)
@@ -103,4 +134,29 @@ async def analyze_wounds(file: UploadFile = File(...)):
         response.headers['X-Area-Cm2'] = str(area_cm2)
         response.headers['X-Moisture'] = str(moisture)
         return response
-    return {'error': 'Invalid file format'}
+
+    return {'error': 'Invalid file format'}
+
+def start_camera():
+    global live_view_running
+    cap = cv2.VideoCapture(0)
+    live_view_running = True
+    while live_view_running:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = draw_square_landmarks(frame)
+        cv2.imshow('Live Landmarks - Press Q to stop', frame)
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            live_view_running = False
+            break
+    cap.release()
+    cv2.destroyAllWindows()
+
+@app.get("/live_landmarks")
+def live_camera_with_landmarks():
+    if not live_view_running:
+        threading.Thread(target=start_camera).start()
+        return {"message": "Live camera started. Check your system's display window."}
+    else:
+        return {"message": "Live camera already running."}