Spaces:
Runtime error
Runtime error
Update modules/tb_image_processor.py
Browse files- modules/tb_image_processor.py +52 -11
modules/tb_image_processor.py
CHANGED
@@ -1,40 +1,81 @@
|
|
|
|
1 |
import cv2
|
2 |
import numpy as np
|
3 |
import logging
|
4 |
from tensorflow.keras.models import load_model
|
5 |
|
|
|
6 |
logging.basicConfig(level=logging.INFO)
|
7 |
logger = logging.getLogger(__name__)
|
8 |
|
9 |
class TBImageProcessor:
|
10 |
-
"""Processes
|
11 |
|
12 |
def __init__(self, model_path="tb_cnn_model.h5"):
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
try:
|
14 |
self.model = load_model(model_path)
|
15 |
-
logger.info("TB Image Processor
|
16 |
except Exception as e:
|
17 |
-
logger.error(f"Failed to load TB Image Model: {e}")
|
18 |
self.model = None
|
19 |
|
20 |
def process_image(self, image_path):
|
21 |
-
"""Analyze
|
22 |
-
|
23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
24 |
|
25 |
try:
|
|
|
26 |
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
27 |
-
|
28 |
-
|
|
|
29 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
30 |
prediction = self.model.predict(image)
|
31 |
confidence = float(prediction[0][0])
|
32 |
result = "TB Detected" if confidence > 0.5 else "No TB"
|
33 |
-
|
|
|
34 |
return {
|
35 |
"result": result,
|
36 |
"confidence": confidence
|
37 |
}
|
|
|
38 |
except Exception as e:
|
39 |
-
logger.error(f"Error
|
40 |
-
return {"error": "
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
import cv2
|
3 |
import numpy as np
|
4 |
import logging
|
5 |
from tensorflow.keras.models import load_model
|
6 |
|
7 |
+
# Set up logging
|
8 |
logging.basicConfig(level=logging.INFO)
|
9 |
logger = logging.getLogger(__name__)
|
10 |
|
11 |
class TBImageProcessor:
|
12 |
+
"""Processes TB images using a trained CNN model for risk assessment."""
|
13 |
|
14 |
def __init__(self, model_path="tb_cnn_model.h5"):
|
15 |
+
# Validate model path
|
16 |
+
if not os.path.exists(model_path):
|
17 |
+
logger.error(f"Model path '{model_path}' does not exist. Please check the path.")
|
18 |
+
self.model = None
|
19 |
+
return
|
20 |
+
|
21 |
try:
|
22 |
self.model = load_model(model_path)
|
23 |
+
logger.info("TB Image Processor model loaded successfully.")
|
24 |
except Exception as e:
|
25 |
+
logger.error(f"Failed to load the TB Image Model: {e}")
|
26 |
self.model = None
|
27 |
|
28 |
def process_image(self, image_path):
|
29 |
+
"""Analyze a TB image and return risk assessment."""
|
30 |
+
# Validate the image file
|
31 |
+
if not os.path.exists(image_path):
|
32 |
+
logger.error(f"Image path '{image_path}' does not exist.")
|
33 |
+
return {"error": "Image file not found."}
|
34 |
+
|
35 |
+
if self.model is None:
|
36 |
+
logger.error("Model is not loaded. Cannot process the image.")
|
37 |
+
return {"error": "Model not loaded."}
|
38 |
|
39 |
try:
|
40 |
+
# Load and preprocess image
|
41 |
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
42 |
+
if image is None:
|
43 |
+
logger.error(f"Failed to read the image at path '{image_path}'.")
|
44 |
+
return {"error": "Invalid image format or corrupted file."}
|
45 |
|
46 |
+
# Resize for CNN input and normalize
|
47 |
+
if image.shape[0] < 128 or image.shape[1] < 128:
|
48 |
+
logger.warning("Image dimensions are smaller than expected, resizing may affect accuracy.")
|
49 |
+
image = cv2.resize(image, (128, 128))
|
50 |
+
image = np.expand_dims(image, axis=[0, -1]) / 255.0
|
51 |
+
|
52 |
+
# Make prediction
|
53 |
prediction = self.model.predict(image)
|
54 |
confidence = float(prediction[0][0])
|
55 |
result = "TB Detected" if confidence > 0.5 else "No TB"
|
56 |
+
|
57 |
+
logger.info(f"Prediction result: {result}, Confidence: {confidence:.2f}")
|
58 |
return {
|
59 |
"result": result,
|
60 |
"confidence": confidence
|
61 |
}
|
62 |
+
|
63 |
except Exception as e:
|
64 |
+
logger.error(f"Error during image processing: {e}")
|
65 |
+
return {"error": f"Failed to process image: {str(e)}"}
|
66 |
+
|
67 |
+
# Example usage
|
68 |
+
if __name__ == "__main__":
|
69 |
+
# Specify the model and image paths
|
70 |
+
model_path = "path/to/your/tb_cnn_model.h5"
|
71 |
+
image_path = "path/to/your/tb_image.jpg"
|
72 |
+
|
73 |
+
# Instantiate the processor and analyze the image
|
74 |
+
processor = TBImageProcessor(model_path=model_path)
|
75 |
+
result = processor.process_image(image_path=image_path)
|
76 |
+
|
77 |
+
# Log or print the final result
|
78 |
+
if "error" in result:
|
79 |
+
logger.error(f"Processing failed: {result['error']}")
|
80 |
+
else:
|
81 |
+
logger.info(f"Final Result: {result['result']}, Confidence: {result['confidence']:.2f}")
|