Upload 8 files

back

@@ -0,0 +1,13 @@
+FROM python:3.11-slim-bullseye
+
+WORKDIR /backend
+
+COPY backend/requirements.txt .
+
+RUN apt-get -y update && pip install --upgrade pip && pip install --no-cache-dir -r requirements.txt && apt-get -y install libgl1 libglib2.0-0
+
+COPY backend .
+
+# ENTRYPOINT ["uvicorn", "main:app", "--port", "5000", "--host", "0.0.0.0", "--reload", "--reload-dir", "/backend"]
+ENTRYPOINT ["python", "main.py", "-f"]
+CMD ["false"]

main.py

@@ -7,16 +7,27 @@ from PIL import Image
 from io import BytesIO
 from model import fetch_model
 from ultralytics import YOLO
+import argparse
+
+parser = argparse.ArgumentParser(description='Smart Vision')
+parser.add_argument('-f','--fetch', help='Do you want to fetch the latest model from the registry? [true, false]', required=False, default=False)
+args = parser.parse_args()
+
+if args.fetch == 'true' or not os.path.isfile('best.pt'):
+    fetch_model()
+model = YOLO('best.pt')
+
+
+# # Load the YOLO model
+# if not os.path.exists(os.path.join('app_temp', 'model', 'best.pt')):
+#     model = fetch_model()  # Load YOLO model
+# else:
+#     model = YOLO(os.path.join('app_temp', 'model', 'best.pt'))
+
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the YOLO model
-if not os.path.exists(os.path.join('app_temp', 'model', 'best.pt')):
-    model = fetch_model()  # Load YOLO model
-else:
-    model = YOLO(os.path.join('app_temp', 'model', 'best.pt'))
-
 def preprocess_image(image_bytes):
     """Convert image bytes to a PIL Image."""
     image = Image.open(BytesIO(image_bytes)).convert("RGB")
@@ -45,68 +56,68 @@ async def predict_image(file: UploadFile = File(...)):
     # Return the image file
     return StreamingResponse(io_buf, media_type="image/jpeg")
     
-@app.post("/process/video")
-async def predict_video(file: UploadFile = File(...)):
-    temp_input = "temp_input.mp4"
-    temp_output = "temp_output.mp4"
+# @app.post("/process/video")
+# async def predict_video(file: UploadFile = File(...)):
+#     temp_input = "temp_input.mp4"
+#     temp_output = "temp_output.mp4"
 
-    with open(temp_input, "wb") as f:
-        f.write(await file.read())
+#     with open(temp_input, "wb") as f:
+#         f.write(await file.read())
 
-    # Open the video
-    cap = cv2.VideoCapture(temp_input)
+#     # Open the video
+#     cap = cv2.VideoCapture(temp_input)
 
-    # Get video properties
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
+#     # Get video properties
+#     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+#     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+#     fps = cap.get(cv2.CAP_PROP_FPS)
 
-    # Create VideoWriter object
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
+#     # Create VideoWriter object
+#     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+#     out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
 
-    # Process each frame
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+#     # Process each frame
+#     while cap.isOpened():
+#         ret, frame = cap.read()
+#         if not ret:
+#             break
 
-        # Convert to RGB for YOLO
-        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+#         # Convert to RGB for YOLO
+#         rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
-        # Run YOLO prediction
-        results = model(rgb_frame)[0]
+#         # Run YOLO prediction
+#         results = model(rgb_frame)[0]
 
-        # Get frame with detection boxes
-        result_frame = results.plot()
+#         # Get frame with detection boxes
+#         result_frame = results.plot()
 
-        # Convert back to BGR for OpenCV
-        result_frame_bgr = cv2.cvtColor(result_frame, cv2.COLOR_RGB2BGR)
+#         # Convert back to BGR for OpenCV
+#         result_frame_bgr = cv2.cvtColor(result_frame, cv2.COLOR_RGB2BGR)
 
-        # Write to output video
-        out.write(result_frame_bgr)
+#         # Write to output video
+#         out.write(result_frame_bgr)
 
-    # Release resources
-    cap.release()
-    out.release()
+#     # Release resources
+#     cap.release()
+#     out.release()
 
-    # Return the video file as response
-    return FileResponse(temp_output, media_type="video/mp4", filename="result.mp4")
+#     # Return the video file as response
+#     return FileResponse(temp_output, media_type="video/mp4", filename="result.mp4")
 
-# @app.post("/process/video")
-# async def predict_video(file: UploadFile = File(...)):
-#     """Predict objects in a video file."""
-#     video_bytes = await file.read()
-#     with open("video.mp4", "wb") as f:
-#         f.write(video_bytes)
+@app.post("/process/video")
+async def predict_video(file: UploadFile = File(...)):
+    """Predict objects in a video file."""
+    video_bytes = await file.read()
+    with open("video.mp4", "wb") as f:
+        f.write(video_bytes)
 
-#     # Perform inference on the video
-#     # model.predict(video_path, save = True, device = 'cuda', project = os.path.join('app_temp', 'results'), name = 'avi', verbose = False, exist_ok = True)
-#     model.predict("video.mp4", device = 'cuda', verbose = False)
+    # Perform inference on the video
+    model.predict('video.mp4', save = True, device = 'cuda', project = 'results', name = 'avi', verbose = False, exist_ok = True)
+    # model.predict("video.mp4", device = 'cuda', verbose = False)
 
-#     # Return the video file
-#     return FileResponse(os.path.join('app_temp', 'result', "vid.avi"), filename="result.mp4")
+    # Return the video file
+    return FileResponse("results/avi/video.avi", filename="result.avi", media_type="video/avi")
 
 
-# if __name__ == "__main__":
-#     uvicorn.run("main:app", port=5000, reload=True)
+if __name__ == "__main__":
+    uvicorn.run("main:app", port=5000, reload=True, host='0.0.0.0')

model.py

@@ -1,8 +1,8 @@
 import mlflow
 from utils import read_yaml, mkdirs
-from os.path import join, exists, dirname
+# from os.path import join, exists, dirname
 from os import environ
-from ultralytics import YOLO
+# from ultralytics import YOLO
 from dotenv import load_dotenv
 import dagshub
 
@@ -12,24 +12,24 @@ dagshub.init(repo_owner='3bdullah3yad', repo_name='SmartVision', mlflow=True)
 
 # model_path = join('app_temp', 'model', 'best.pt')
 
-def fetch_model(repull=False):
-    if not exists('best.pt') or repull:
-        # mkdirs([model_path])
-        params = read_yaml("params.yaml")
-        cfg = read_yaml("config.yaml")
-        
-        if params.inference.exp_name == 'from_cfg_train':
-            experiment = dict(mlflow.get_experiment_by_name(cfg.train.exp_name))
-        else:
-            experiment = dict(mlflow.get_experiment_by_name(params.inference.exp_name))
+def fetch_model():
+    # mkdirs([model_path])
+    params = read_yaml("params.yaml")
+    cfg = read_yaml("config.yaml")
+    
+    if params.inference.exp_name == 'from_cfg_train':
+        experiment = dict(mlflow.get_experiment_by_name(cfg.train.exp_name))
+    else:
+        experiment = dict(mlflow.get_experiment_by_name(params.inference.exp_name))
 
-        run_id = params.inference.run_id
-        if run_id == 'latest':
-            experiment_id = experiment['experiment_id']
-            df = mlflow.search_runs([experiment_id], order_by=["end_time DESC"])
-            run_id = df['run_id'][0]        
+    run_id = params.inference.run_id
+    if run_id == 'latest':
+        experiment_id = experiment['experiment_id']
+        df = mlflow.search_runs([experiment_id], order_by=["end_time DESC"])
+        run_id = df['run_id'][0]        
 
-        # mlflow.artifacts.download_artifacts(f'runs:/{run_id}/weights/best.pt', dst_path = dirname(model_path))
-        mlflow.artifacts.download_artifacts(f'runs:/{run_id}/weights/best.pt', dst_path = dirname('best.pt'))
+    # mlflow.artifacts.download_artifacts(f'runs:/{run_id}/weights/best.pt', dst_path = dirname(model_path))
+    mlflow.artifacts.download_artifacts(f'runs:/{run_id}/weights/best.pt', dst_path = '.')
+    print(f"\n\nModel downloaded from MLFlow, Run ID: {run_id}\n\n")
 
-    return YOLO('best.pt')
+    # return YOLO('best.pt')