Fixing model and most of UI

models/{model.onnx → model.ckpt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1e564444ab1d45dd415b083ae5478368717b508748c6653ba56b457574db619c
-size 85128734
+oid sha256:5601b234e608862cde6159ba32bd77a3e5e2b23e41ce488ee778bf4154419090
+size 94409193

requirements.txt

@@ -3,5 +3,4 @@ torch
 torchvision
 pytorch-lightning
 matplotlib
-numpy
-onnxruntime
+numpy

app.py → scripts/app.py

@@ -4,50 +4,103 @@ from PIL import Image
 import pytorch_lightning as pl
 import torch.nn as nn
 from torchvision import transforms as T
-from torchvision.models import resnet152
+from torchvision import models
 import matplotlib.pyplot as plt
 import onnxruntime as ort
 from glob import glob
 import streamlit as st
 import numpy as np
+from torchmetrics.functional import accuracy
+from torchmetrics import Accuracy
 
 #Define the labels
-labels = ['Defect', 'nDefect']
+labels = ['Defect', 'Non-Defect']
 
 # Define the sample images
 sample_images = {
-    "Defect01": "pics/Defect/0.jpg",
-    "Defect02": "pics/Defect/1.jpg",
-    "Defect03": "pics/Defect/2.jpg",
-    "Non-Defect01": "pics/nDefect/0.jpg",
-    "Non-Defect02": "pics/nDefect/1.jpg",
-    "Non-Defect03": "pics/nDefect/2.jpg"
+    "Defect01": "../pics/Defect/2.jpg",
+    "Defect02": "../pics/Defect/6.jpg",
+    "Defect03": "../pics/Defect/8.jpg",
+    "Non-Defect01": "../pics/nDefect/3.jpg",
+    "Non-Defect02": "../pics/nDefect/4.jpg",
+    "Non-Defect03": "../pics/nDefect/8.jpg"
 }
 
-class ONNXModel(pl.LightningModule):
-    def __init__(self, model_path, n_classes=2):
-        super().__init__()
-        self.onnx_session = ort.InferenceSession(model_path)
-        self.input_name = self.onnx_session.get_inputs()[0].name
+class DefectResNet(pl.LightningModule):
+    def __init__(self, n_classes=2):
+        super(DefectResNet, self).__init__()
+        
+        # จำนวนของพันธุ์output (2)
         self.n_classes = n_classes
-        #Change the last layer
-        self.backbone = resnet152(weights="ResNet152_Weights.DEFAULT")
+
+        #เปลี่ยน layer สุดท้าย
+        self.backbone = models.resnet50(pretrained=True)
         # self.backbone = models.resnet152(pretrained=True)
         # self.backbone = models.vgg19(pretrained=True)
         for param in self.backbone.parameters():
             param.requires_grad = False
-        #Change the last layer to 2 classes
+
+        # เปลี่ยน fc layer เป็น output ขนาด 2
         self.backbone.fc = torch.nn.Linear(self.backbone.fc.in_features, n_classes) #For ResNet base mdoel
         # self.backbone.classifier[6] = torch.nn.Linear(self.backbone.classifier[6].in_features, n_classes) #For VGG bse model
-        self.onnx_session = ort.InferenceSession(model_path)
-        self.input_name = self.onnx_session.get_inputs()[0].name
+        
+        self.entropy_loss = nn.CrossEntropyLoss()
+        self.accuracy = Accuracy(task="multiclass", num_classes=2)
+
+        self.save_hyperparameters(logger=False)
+
     def forward(self, x):
         preds = self.backbone(x)
         return preds
 
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.backbone(x)
+        loss = self.entropy_loss(logits, y)
+        y_pred = torch.argmax(logits, dim=1)
+        self.log("train_loss", loss)
+        self.log("train_acc", self.accuracy(y_pred, y))
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.backbone(x)
+        loss = self.entropy_loss(logits, y)
+        y_pred = torch.argmax(logits, dim=1)
+        self.log("val_loss", loss)
+        self.log("val_acc", self.accuracy(y_pred, y))
+        return loss
+        
+    def configure_optimizers(self):
+        self.optimizer = torch.optim.AdamW(self.parameters(), lr=1e-3)
+        return {
+            "optimizer": self.optimizer,
+            "monitor": "val_loss",
+        }
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.backbone(x)
+        loss = self.entropy_loss(logits, y)
+        y_pred = torch.argmax(logits, dim=1)
+        self.log("val_loss", loss)
+        self.log("val_acc", self.accuracy(y_pred, y))
+        return loss
+    
+    def _shared_eval_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.model(x)
+        logits = self.backbone(x)
+        loss = self.entropy_loss(logits, y)
+        acc = accuracy(y_hat, y)
+        return loss, acc
+
 # Load the model on the appropriate device
-model = ONNXModel(model_path='models/model.onnx')
-model = model.to("cpu")
+loadmodel = DefectResNet()
+def load_checkpoint(checkpoint):
+    loadmodel.load_state_dict(checkpoint["state_dict"])
+load_checkpoint(torch.load("../models/model.ckpt"))
+loadmodel.eval()
 
 transform = T.Compose([
     T.Resize((224, 224)),
@@ -56,24 +109,21 @@ transform = T.Compose([
 
 def predict(image):
     image = transform(image).unsqueeze(0)
-    image = image.to("cpu")
 
     # Perform the prediction
     with torch.no_grad():
-        logits = model(image)
+        logits = loadmodel(image)
         probs = F.softmax(logits, dim=1)
         return probs
 
 # Define the Streamlit app
 def app():
     predictions = None
-    st.title("Defect Classification")
-
+    st.title("Digital textile printing defect classification for industrial.")
     uploaded_file = st.file_uploader("Upload your image...", type=["jpg"])
 
     with st.expander("Or choose from sample here..."):
-        
-        st.header("Sample Defect Images")
+        sample = st.selectbox(label = "Select here", options = list(sample_images.keys()), label_visibility="hidden")   
         col1, col2, col3 = st.columns(3)
         with col1:
             st.image(sample_images["Defect01"], caption="Defect01", use_column_width=True)
@@ -81,7 +131,6 @@ def app():
             st.image(sample_images["Defect02"], caption="Defect02", use_column_width=True)
         with col3:
             st.image(sample_images["Defect03"], caption="Defect03", use_column_width=True)
-        st.header("Sample Non-Defect Images")
         col1, col2, col3 = st.columns(3)
         with col1:
             st.image(sample_images["Non-Defect01"], caption="Non-Defect01", use_column_width=True)
@@ -90,8 +139,6 @@ def app():
         with col3:
             st.image(sample_images["Non-Defect03"], caption="Non-Defect03", use_column_width=True)
 
-        sample = st.selectbox(label = "Select here", options = list(sample_images.keys()), label_visibility="hidden")   
-
     # If an image is uploaded, make a prediction on it
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
@@ -104,13 +151,16 @@ def app():
       
     # Show  predictions with their probabilities
     if predictions is not  None:
-        st.write(predictions)
+        # st.write(predictions)
+        st.subheader(f'Predictions : {labels[torch.argmax(predictions[0]).item()]}')
         for pred, prob in zip(labels, predictions[0]):
             st.write(f"{pred}: {prob * 100:.2f}%")
             st.progress(prob.item())
     else:
         st.write("No predictions.")
-
+    st.subheader("Credits")
+    st.write("By : Settapun Laoaree | AI-Builders")
+    st.markdown("Source : [Github](https://github.com/ShokulSet/DefectDetection-AIBuilders) [Hugging Face](https://huggingface.co/spaces/sh0kul/DefectDetection-Deploy)")
 
 # Run the app
 if __name__ == "__main__":

pics/Rename.py → scripts/rename.py

@@ -1,5 +1,5 @@
 # Python 3 code to rename multiple
-# files in a directory or folder
+# files in a directory or defect_path
  
 # importing os module
 import os
@@ -7,15 +7,24 @@ import os
 # Function to rename multiple files
 def main():
    
-    folder = "/home/shokul/AIBuilder/deploy/DefectDetection-AIBuilders/pics/nDefect"
-    for count, filename in enumerate(os.listdir(folder)):
+    defect_path = "../pics/Defect"
+    ndefect_path = "../pics/nDefect"
+    for count, filename in enumerate(os.listdir(defect_path)):
         dst = f"{str(count)}.jpg"
-        src =f"{folder}/{filename}"  # foldername/filename, if .py file is outside folder
-        dst =f"{folder}/{dst}"
+        src =f"{defect_path}/{filename}"  # defect_pathname/filename, if .py file is outside defect_path
+        dst =f"{defect_path}/{dst}"
          
         # rename() function will
         # rename all the files
         os.rename(src, dst)
+
+    for count, filename in enumerate(os.listdir(ndefect_path)):
+        dst = f"{str(count)}.jpg"
+        src =f"{ndefect_path}/{filename}"  # defect_pathname/filename, if .py file is outside defect_path
+        dst =f"{ndefect_path}/{dst}"
+        # rename() function will
+        # rename all the files
+        os.rename(src, dst)
  
 # Driver Code
 if __name__ == '__main__':