First

app.py

@@ -0,0 +1,167 @@
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import pytorch_lightning as pl
+import torch.nn as nn
+from torchvision import transforms as T
+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', 'Non-Defect']
+
+# Define the sample images
+sample_images = {
+    "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 DefectResNet(pl.LightningModule):
+    def __init__(self, n_classes=2):
+        super(DefectResNet, self).__init__()
+        
+        # จำนวนของพันธุ์output (2)
+        self.n_classes = n_classes
+
+        #เปลี่ยน 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
+
+        # เปลี่ยน 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.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
+loadmodel = DefectResNet()
+def load_checkpoint(checkpoint):
+    loadmodel.load_state_dict(checkpoint["state_dict"])
+load_checkpoint(torch.load("models/model.ckpt", map_location=torch.device('cpu')))
+loadmodel.eval()
+
+transform = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor()
+])
+
+def predict(image):
+    image = transform(image).unsqueeze(0)
+
+    # Perform the prediction
+    with torch.no_grad():
+        logits = loadmodel(image)
+        probs = F.softmax(logits, dim=1)
+        return probs
+
+# Define the Streamlit app
+def app():
+    predictions = None
+    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..."):
+        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)
+        with col2:
+            st.image(sample_images["Defect02"], caption="Defect02", use_column_width=True)
+        with col3:
+            st.image(sample_images["Defect03"], caption="Defect03", use_column_width=True)
+        col1, col2, col3 = st.columns(3)
+        with col1:
+            st.image(sample_images["Non-Defect01"], caption="Non-Defect01", use_column_width=True)
+        with col2:
+            st.image(sample_images["Non-Defect02"], caption="Non-Defect02", use_column_width=True)
+        with col3:
+            st.image(sample_images["Non-Defect03"], caption="Non-Defect03", use_column_width=True)
+
+    # If an image is uploaded, make a prediction on it
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+        predictions = predict(image)
+    elif sample:    
+        image = Image.open(sample_images[sample])
+        st.image(image, caption=sample.capitalize() + " Image", use_column_width=True)
+        predictions = predict(image)
+      
+    # Show  predictions with their probabilities
+    if predictions is not  None:
+        # 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__":
+    app()

models/model.ckpt

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

requirements.txt

@@ -0,0 +1,89 @@
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==4.2.2
+async-timeout==4.0.2
+attrs==23.1.0
+blinker==1.6.2
+cachetools==5.3.0
+certifi==2023.5.7
+charset-normalizer==3.1.0
+click==8.1.3
+cmake==3.26.3
+coloredlogs==15.0.1
+contourpy==1.0.7
+cycler==0.11.0
+decorator==5.1.1
+entrypoints==0.4
+filelock==3.12.0
+flatbuffers==23.5.9
+fonttools==4.39.4
+frozenlist==1.3.3
+fsspec==2023.5.0
+gitdb==4.0.10
+GitPython==3.1.31
+humanfriendly==10.0
+idna==3.4
+importlib-metadata==6.6.0
+Jinja2==3.1.2
+jsonschema==4.17.3
+kiwisolver==1.4.4
+lightning-utilities==0.8.0
+lit==16.0.5
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.1
+numpy==1.24.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+onnxruntime==1.14.1
+packaging==23.1
+pandas==2.0.1
+Pillow==9.5.0
+protobuf==3.20.3
+pyarrow==12.0.0
+pydeck==0.8.1b0
+Pygments==2.15.1
+Pympler==1.0.1
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+pytorch-lightning==2.0.2
+pytz==2023.3
+pytz-deprecation-shim==0.1.0.post0
+PyYAML==6.0
+requests==2.31.0
+rich==13.3.5
+six==1.16.0
+smmap==5.0.0
+streamlit==1.23.1
+sympy==1.12
+tenacity==8.2.2
+toml==0.10.2
+toolz==0.12.0
+torch==2.0.1
+torchmetrics==0.11.4
+torchvision==0.15.2
+tornado==6.3.2
+tqdm==4.65.0
+triton==2.0.0
+typing_extensions==4.6.0
+tzdata==2023.3
+tzlocal==4.3
+urllib3==2.0.2
+validators==0.20.0
+watchdog==3.0.0
+yarl==1.9.2
+zipp==3.15.0

scripts/rename.py

@@ -0,0 +1,33 @@
+# Python 3 code to rename multiple
+# files in a directory or defect_path
+ 
+# importing os module
+import os
+ 
+# Function to rename multiple files
+def main():
+   
+    defect_path = "../pics/Defect"
+    ndefect_path = "../pics/nDefect"
+    for count, filename in enumerate(os.listdir(defect_path)):
+        dst = f"{str(count)}.jpg"
+        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__':
+     
+    # Calling main() function
+    main()