Raiff1982 commited on
Commit
a6e5aa5
·
verified ·
1 Parent(s): c7f72d8

Update modules/tb_image_processor.py

Browse files
Files changed (1) hide show
  1. 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 real TB saliva images using a trained CNN model"""
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 Model Loaded Successfully.")
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 the TB image and return risk assessment."""
22
- if not self.model:
23
- return {"error": "Model not loaded. Cannot process image."}
 
 
 
 
 
 
24
 
25
  try:
 
26
  image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
27
- image = cv2.resize(image, (128, 128)) # Resizing for CNN input
28
- image = np.expand_dims(image, axis=[0, -1]) / 255.0 # Normalize
 
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 processing image: {e}")
40
- return {"error": "Image processing failed."}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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}")