tb update

Dockerfile

@@ -15,6 +15,7 @@ RUN pip install torch torchvision --index-url https://download.pytorch.org/whl/c
 RUN pip install torchxrayvision
 RUN pip install grad-cam
 RUN pip install opencv-python
+RUN pip install transformers
 
 COPY --chown=user . /app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -5,9 +5,11 @@ import cv2
 import numpy as np
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-
-
+from transformers import pipeline
+from PIL import Image
 from fastapi.middleware.cors import CORSMiddleware
+import requests
+
 app = FastAPI()
 # Add the frontend origin here
 origins = [
@@ -25,7 +27,7 @@ app.add_middleware(
 
 
 model = xrv.models.DenseNet(weights="densenet121-res224-all")
-
+tb_classifier =pipeline("image-classification",model="vimal-humantics/dinov2-base-xray-224-finetuned-tb") 
 
 
 
@@ -82,8 +84,20 @@ def predict(image_url:str = Query(..., description="URL to a chest X-ray image")
             pred_output.update({k:round(v,2)})
 
         get_bounding_box = show_anomaly_bounding_box(img,model=model)
-
-        return {"prediction_result":pred_output,"bounding_box":{pred_label:get_bounding_box}}
+        # TB detection
+        
+        image = Image.open(requests.get(image_url, stream=True).raw)
+        tb_finding = tb_classifier(images=image)
+        tb_label = tb_finding[0]['label']
+        print(tb_label)
+        tb_score = round(tb_finding[0]['score'],2)
+        tb_output = 0
+        if tb_label == "normal":
+            tb_output = 1-tb_score
+        else:
+            tb_output = tb_score
+
+        return {"prediction_result":pred_output,"bounding_box":{pred_label:get_bounding_box},"tb_finding":tb_output}
     except Exception as e:
         print(e)
         raise HTTPException(status_code=400, detail=f"Failed to fetch/process image: {str(e)}")