<<<<<<< HEAD
import os
import cv2
import numpy as np
from flask import Flask, render_template, request, jsonify, send_from_directory
import torch
from torchvision import transforms, models
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image
import uuid


app = Flask(__name__)

# Class names
class_names = ['Covid-19', 'Emphysema', 'Healthy', 'Pneumonia', 'Tuberculosis', 'Random']

# Load the pre-trained ResNet101 model
model = models.resnet101(pretrained=False)
model.fc = torch.nn.Linear(in_features=2048, out_features=len(class_names))
checkpoint_path = "resnet101_state_dict.pth"  
model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
model.eval()

# Image transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Function to preprocess image
def preprocess_image(image_path):
    image = Image.open(image_path).convert('RGB')
    return transform(image).unsqueeze(0)

# Generate visualizations for Grad-CAM and other images
def generate_visualizations(image_path):
    original_image = cv2.imread(image_path)
    original_image_rgb = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)

    # Grayscale image
    grayscale_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)

    # Histogram equalized image
    equalized_image = cv2.equalizeHist(grayscale_image)

    # Edge detection result
    edges_image = cv2.Canny(grayscale_image, 50, 150)

    # Segmented image
    _, segmented_image = cv2.threshold(grayscale_image, 127, 255, cv2.THRESH_BINARY)

    # Grad-CAM Visualization
    input_tensor = preprocess_image(image_path)
    target_layer = model.layer4[-1]  # Last layer of ResNet101
    cam = GradCAM(model=model, target_layers=[target_layer])
    grayscale_cam = cam(input_tensor=input_tensor)[0]

    # Normalize the Grad-CAM output to range [0, 1]
    grayscale_cam = np.maximum(grayscale_cam, 0)
    grayscale_cam = grayscale_cam / np.max(grayscale_cam)

    # Apply the Grad-CAM heatmap on the image with a red color map
    input_image_rgb_resized = cv2.resize(original_image_rgb, (224, 224))

    # Using a custom colormap (ensure red regions are highlighted)
    heatmap = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
    grad_cam_image = cv2.addWeighted(input_image_rgb_resized, 0.7, heatmap, 0.3, 0)

    # Mask for ROI extraction (high-confidence areas are those with grayscale_cam > 0.5)
    roi_mask = (grayscale_cam > 0.5).astype(np.uint8)
    roi = cv2.bitwise_and(input_image_rgb_resized, input_image_rgb_resized, mask=roi_mask)

    # Generate unique file name
    file_id = str(uuid.uuid4())

    # Save images to server and return file paths
    visualization_paths = {
        "original": f'/uploads/{file_id}_original.png',
        "grayscale": f'/uploads/{file_id}_grayscale.png',
        "equalized": f'/uploads/{file_id}_equalized.png',
        "edges": f'/uploads/{file_id}_edges.png',
        "segmented": f'/uploads/{file_id}_segmented.png',
        "grad_cam": f'/uploads/{file_id}_grad_cam.png',
        "roi": f'/uploads/{file_id}_roi.png'
    }

    # Create uploads directory if it doesn't exist
    os.makedirs('uploads', exist_ok=True)

    # Save the visualizations as images
    cv2.imwrite(f'uploads/{file_id}_original.png', original_image)
    cv2.imwrite(f'uploads/{file_id}_grayscale.png', grayscale_image)
    cv2.imwrite(f'uploads/{file_id}_equalized.png', equalized_image)
    cv2.imwrite(f'uploads/{file_id}_edges.png', edges_image)
    cv2.imwrite(f'uploads/{file_id}_segmented.png', segmented_image)
    cv2.imwrite(f'uploads/{file_id}_grad_cam.png', grad_cam_image)
    cv2.imwrite(f'uploads/{file_id}_roi.png', roi)

    return visualization_paths

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    image_id = str(uuid.uuid4())
    file_path = os.path.join('uploads', f'{image_id}.png')
    os.makedirs('uploads', exist_ok=True)
    file.save(file_path)

    input_tensor = preprocess_image(file_path)
    with torch.no_grad():
        output = model(input_tensor)
        predicted_class = output.argmax(dim=1).item()
        confidence_score = torch.nn.functional.softmax(output, dim=1)[0, predicted_class].item() * 100

    result = {
        "predicted_class": class_names[predicted_class],
        "confidence_score": f"{confidence_score:.2f}%",
    }

    if class_names[predicted_class] == "Random":
        result["message"] = "Sorry... You inserted a Non X-ray image. Please try again with a chest x-ray image. Thank you."
        result["visualizations"] = None
    else:
        visualizations = generate_visualizations(file_path)
        result["visualizations"] = visualizations

    return jsonify(result)

# To serve the images from the 'uploads' folder correctly
@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory('uploads', filename)

if __name__ == '__main__':
    app.run(debug=True)
=======
import os
import cv2
import numpy as np
from flask import Flask, render_template, request, jsonify, send_from_directory
import torch
from torchvision import transforms, models
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image
import uuid


app = Flask(__name__)

# Class names
class_names = ['Covid-19', 'Emphysema', 'Healthy', 'Pneumonia', 'Tuberculosis', 'Random']

# Load the pre-trained ResNet101 model
model = models.resnet101(pretrained=False)
model.fc = torch.nn.Linear(in_features=2048, out_features=len(class_names))
checkpoint_path = "resnet101_state_dict.pth"  
model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))
model.eval()

# Image transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Function to preprocess image
def preprocess_image(image_path):
    image = Image.open(image_path).convert('RGB')
    return transform(image).unsqueeze(0)

# Generate visualizations for Grad-CAM and other images
def generate_visualizations(image_path):
    original_image = cv2.imread(image_path)
    original_image_rgb = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)

    # Grayscale image
    grayscale_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)

    # Histogram equalized image
    equalized_image = cv2.equalizeHist(grayscale_image)

    # Edge detection result
    edges_image = cv2.Canny(grayscale_image, 50, 150)

    # Segmented image
    _, segmented_image = cv2.threshold(grayscale_image, 127, 255, cv2.THRESH_BINARY)

    # Grad-CAM Visualization
    input_tensor = preprocess_image(image_path)
    target_layer = model.layer4[-1]  # Last layer of ResNet101
    cam = GradCAM(model=model, target_layers=[target_layer])
    grayscale_cam = cam(input_tensor=input_tensor)[0]

    # Normalize the Grad-CAM output to range [0, 1]
    grayscale_cam = np.maximum(grayscale_cam, 0)
    grayscale_cam = grayscale_cam / np.max(grayscale_cam)

    # Apply the Grad-CAM heatmap on the image with a red color map
    input_image_rgb_resized = cv2.resize(original_image_rgb, (224, 224))

    # Using a custom colormap (ensure red regions are highlighted)
    heatmap = cv2.applyColorMap(np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET)
    grad_cam_image = cv2.addWeighted(input_image_rgb_resized, 0.7, heatmap, 0.3, 0)

    # Mask for ROI extraction (high-confidence areas are those with grayscale_cam > 0.5)
    roi_mask = (grayscale_cam > 0.5).astype(np.uint8)
    roi = cv2.bitwise_and(input_image_rgb_resized, input_image_rgb_resized, mask=roi_mask)

    # Generate unique file name
    file_id = str(uuid.uuid4())

    # Save images to server and return file paths
    visualization_paths = {
        "original": f'/uploads/{file_id}_original.png',
        "grayscale": f'/uploads/{file_id}_grayscale.png',
        "equalized": f'/uploads/{file_id}_equalized.png',
        "edges": f'/uploads/{file_id}_edges.png',
        "segmented": f'/uploads/{file_id}_segmented.png',
        "grad_cam": f'/uploads/{file_id}_grad_cam.png',
        "roi": f'/uploads/{file_id}_roi.png'
    }

    # Create uploads directory if it doesn't exist
    os.makedirs('uploads', exist_ok=True)

    # Save the visualizations as images
    cv2.imwrite(f'uploads/{file_id}_original.png', original_image)
    cv2.imwrite(f'uploads/{file_id}_grayscale.png', grayscale_image)
    cv2.imwrite(f'uploads/{file_id}_equalized.png', equalized_image)
    cv2.imwrite(f'uploads/{file_id}_edges.png', edges_image)
    cv2.imwrite(f'uploads/{file_id}_segmented.png', segmented_image)
    cv2.imwrite(f'uploads/{file_id}_grad_cam.png', grad_cam_image)
    cv2.imwrite(f'uploads/{file_id}_roi.png', roi)

    return visualization_paths

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    image_id = str(uuid.uuid4())
    file_path = os.path.join('uploads', f'{image_id}.png')
    os.makedirs('uploads', exist_ok=True)
    file.save(file_path)

    input_tensor = preprocess_image(file_path)
    with torch.no_grad():
        output = model(input_tensor)
        predicted_class = output.argmax(dim=1).item()
        confidence_score = torch.nn.functional.softmax(output, dim=1)[0, predicted_class].item() * 100

    result = {
        "predicted_class": class_names[predicted_class],
        "confidence_score": f"{confidence_score:.2f}%",
    }

    if class_names[predicted_class] == "Random":
        result["message"] = "Sorry... You inserted a Non X-ray image. Please try again with a chest x-ray image. Thank you."
        result["visualizations"] = None
    else:
        visualizations = generate_visualizations(file_path)
        result["visualizations"] = visualizations

    return jsonify(result)

# To serve the images from the 'uploads' folder correctly
@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory('uploads', filename)

if __name__ == '__main__':
    app.run(debug=True)
>>>>>>> 5676c55 (Add ResNet101 model using Git LFS)