junaid17 commited on
Commit
f80a3ac
·
verified ·
1 Parent(s): 60d8a59

Upload 4 files

Browse files
Files changed (4) hide show
  1. Dockerfile +26 -0
  2. README.md +4 -4
  3. app.py +139 -0
  4. requirements.txt +12 -0
Dockerfile ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.10-slim
2
+
3
+ ENV PYTHONDONTWRITEBYTECODE=1
4
+ ENV PYTHONUNBUFFERED=1
5
+
6
+ WORKDIR /app
7
+
8
+ # --- SYSTEM DEPENDENCIES (CRITICAL FOR OPENCV / YOLO) ---
9
+ RUN apt-get update && apt-get install -y \
10
+ build-essential \
11
+ gcc \
12
+ libgl1 \
13
+ libglib2.0-0 \
14
+ && rm -rf /var/lib/apt/lists/*
15
+
16
+ # --- PYTHON DEPENDENCIES ---
17
+ COPY requirements.txt .
18
+ RUN pip install --no-cache-dir --upgrade pip \
19
+ && pip install --no-cache-dir -r requirements.txt
20
+
21
+ # --- APP CODE ---
22
+ COPY . .
23
+
24
+ EXPOSE 7860
25
+
26
+ CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
README.md CHANGED
@@ -1,8 +1,8 @@
1
  ---
2
- title: DamageLens
3
- emoji: 🔥
4
- colorFrom: blue
5
- colorTo: blue
6
  sdk: docker
7
  pinned: false
8
  ---
 
1
  ---
2
+ title: DamageLensAI
3
+ emoji: 😻
4
+ colorFrom: red
5
+ colorTo: green
6
  sdk: docker
7
  pinned: false
8
  ---
app.py ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import uuid
3
+ import shutil
4
+ from fastapi import FastAPI, UploadFile, File, HTTPException
5
+ from fastapi.staticfiles import StaticFiles
6
+ from PIL import Image
7
+ from fastapi.middleware.cors import CORSMiddleware
8
+ from scripts.gradcam import get_resnet_gradcam, get_deit_gradcam
9
+ from scripts.yolo import get_yolo_damage_boxes
10
+ from scripts.prediction_helper import ResnetCarDamagePredictor, DeitCarDamagePredictor, FusionCarDamagePredictor
11
+
12
+ app = FastAPI()
13
+
14
+ app.add_middleware(
15
+ CORSMiddleware,
16
+ allow_origins=["*"],
17
+ allow_credentials=True,
18
+ allow_methods=["*"],
19
+ allow_headers=["*"],
20
+ )
21
+
22
+ UPLOAD_DIR = "static/uploads"
23
+ RESULT_DIR = "static/results"
24
+ os.makedirs(UPLOAD_DIR, exist_ok=True)
25
+ os.makedirs(RESULT_DIR, exist_ok=True)
26
+
27
+ app.mount("/static", StaticFiles(directory="static"), name="static")
28
+
29
+ class_map = {
30
+ 0: "Front Breakage",
31
+ 1: "Front Crushed",
32
+ 2: "Front Normal",
33
+ 3: "Rear Breakage",
34
+ 4: "Rear Crushed",
35
+ 5: "Rear Normal"
36
+ }
37
+
38
+ resnet_checkpoint = "checkpoints/best_resnet_model.pt"
39
+ deit_checkpoint = "checkpoints/best_deit_model.pt"
40
+
41
+
42
+ Resnet_Model = ResnetCarDamagePredictor(resnet_checkpoint, class_map)
43
+ Deit_Model = DeitCarDamagePredictor(deit_checkpoint, class_map)
44
+ Fusion_Model = FusionCarDamagePredictor(resnet_predictor=Resnet_Model, deit_predictor=Deit_Model, resnet_weight=0.5, deit_weight=0.5)
45
+
46
+ resnet_predictor = Resnet_Model
47
+ deit_predictor = Deit_Model
48
+
49
+ # ====================== API Endpoint ======================
50
+
51
+ @app.get("/")
52
+ def api_status():
53
+ return {"status": "API is running"}
54
+
55
+ # ============================= Grad-CAM Generation Endpoint =============================
56
+
57
+ @app.post("/predict")
58
+ async def predict_and_generate_cams(file: UploadFile = File(...)):
59
+ unique_id = str(uuid.uuid4())
60
+ input_filename = f"{unique_id}_input.jpg"
61
+ resnet_out_name = f"{unique_id}_resnet.jpg"
62
+ deit_out_name = f"{unique_id}_deit.jpg"
63
+
64
+ input_path = os.path.join(UPLOAD_DIR, input_filename)
65
+ resnet_path = os.path.join(RESULT_DIR, resnet_out_name)
66
+ deit_path = os.path.join(RESULT_DIR, deit_out_name)
67
+
68
+ # Save uploaded file
69
+ with open(input_path, "wb") as buffer:
70
+ shutil.copyfileobj(file.file, buffer)
71
+
72
+ # Generate Grad-CAMs
73
+ get_resnet_gradcam(input_path, resnet_predictor, resnet_path)
74
+ get_deit_gradcam(input_path, deit_predictor, deit_path)
75
+
76
+ # Return the URLs
77
+ return {
78
+ "status": "success",
79
+ "original_image": f"/static/uploads/{input_filename}",
80
+ "resnet_viz": f"/static/results/{resnet_out_name}",
81
+ "deit_viz": f"/static/results/{deit_out_name}"
82
+ }
83
+
84
+ # ============================= Prediction-Only Endpoints =============================
85
+ # ============================= Resnet Prediction =====================================
86
+
87
+ @app.post("/predict/resnet")
88
+ async def resnet_prediction(image : UploadFile = File(...)):
89
+ try:
90
+ image = Image.open(image.file)
91
+ except Exception:
92
+ raise HTTPException(status_code=400, detail="Invalid image file")
93
+ result = Resnet_Model.resnet_predict(image_input=image)
94
+ return result
95
+
96
+ # ============================= Deit Prediction =====================================
97
+ @app.post("/predict/deit")
98
+ async def deit_prediction(image : UploadFile = File(...)):
99
+ try:
100
+ image = Image.open(image.file)
101
+ except Exception:
102
+ raise HTTPException(status_code=400, detail="Invalid image file")
103
+ result = Deit_Model.deit_predict(image_input=image)
104
+ return result
105
+
106
+ # ============================= Fusion Prediction =====================================
107
+ @app.post("/predict/fusion")
108
+ async def fusion_prediction(image : UploadFile = File(...)):
109
+ try:
110
+ image = Image.open(image.file)
111
+ except Exception:
112
+ raise HTTPException(status_code=400, detail="Invalid image file")
113
+ result = Fusion_Model.fuse_predict(image_input=image)
114
+ return result
115
+
116
+ # ============================= YOLO Damage Box Endpoint =============================
117
+ @app.post("/predict/yolo")
118
+ async def yolo_detection(file: UploadFile = File(...)):
119
+ unique_id = str(uuid.uuid4())
120
+
121
+ input_filename = f"{unique_id}_input.jpg"
122
+ yolo_out_name = f"{unique_id}_yolo.jpg"
123
+
124
+ input_path = os.path.join(UPLOAD_DIR, input_filename)
125
+ yolo_path = os.path.join(RESULT_DIR, yolo_out_name)
126
+
127
+ with open(input_path, "wb") as buffer:
128
+ shutil.copyfileobj(file.file, buffer)
129
+
130
+ result = get_yolo_damage_boxes(input_path, yolo_path)
131
+
132
+ return {
133
+ "status": "success",
134
+ "original_image": f"/static/uploads/{input_filename}",
135
+ "yolo_image": f"/static/results/{yolo_out_name}",
136
+ "detections": result["detections"],
137
+ "total_detections": result["total_detections"],
138
+ "message": result["message"]
139
+ }
requirements.txt ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ transformers
4
+ fastapi
5
+ uvicorn
6
+ dotenv
7
+ matplotlib
8
+ opencv-python
9
+ python-multipart
10
+ ultralytics
11
+ plotly
12
+ pandas