Raiff1982 commited on
Commit
f23140a
·
verified ·
1 Parent(s): 8dd90df

Create tb_image_processor.py

Browse files
Files changed (1) hide show
  1. modules/tb_image_processor.py +41 -0
modules/tb_image_processor.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import cv2
3
+ import numpy as np
4
+ import logging
5
+ from tensorflow.keras.models import load_model
6
+
7
+ logging.basicConfig(level=logging.INFO)
8
+ logger = logging.getLogger(__name__)
9
+
10
+ class TBImageProcessor:
11
+ """Processes real TB saliva images using a trained CNN model"""
12
+
13
+ def __init__(self, model_path="tb_cnn_model.h5"):
14
+ try:
15
+ self.model = load_model(model_path)
16
+ logger.info("TB Image Processor Model Loaded Successfully.")
17
+ except Exception as e:
18
+ logger.error(f"Failed to load TB Image Model: {e}")
19
+ self.model = None
20
+
21
+ def process_image(self, image_path):
22
+ """Analyze the TB image and return risk assessment."""
23
+ if not self.model:
24
+ return {"error": "Model not loaded. Cannot process image."}
25
+
26
+ try:
27
+ image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
28
+ image = cv2.resize(image, (128, 128)) # Resizing for CNN input
29
+ image = np.expand_dims(image, axis=[0, -1]) / 255.0 # Normalize
30
+
31
+ prediction = self.model.predict(image)
32
+ confidence = float(prediction[0][0])
33
+ result = "TB Detected" if confidence > 0.5 else "No TB"
34
+
35
+ return {
36
+ "result": result,
37
+ "confidence": confidence
38
+ }
39
+ except Exception as e:
40
+ logger.error(f"Error processing image: {e}")
41
+ return {"error": "Image processing failed."}