File size: 3,332 Bytes
39cdc68 4cb22d8 39cdc68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
import numpy as np
import torch
from skorch import NeuralNetClassifier
from torch import nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
mnist = fetch_openml('mnist_784', as_frame=False, cache=False)
X = mnist.data.astype('float32')
y = mnist.target.astype('int64')
X /= 255.0
device = 'cuda' if torch.cuda.is_available() else 'cpu'
XCnn = X.reshape(-1, 1, 28, 28)
XCnn_train, XCnn_test, y_train, y_test = train_test_split(XCnn, y, test_size=0.25, random_state=42)
from PIL import Image
import torchvision.transforms as transforms
class Cnn(nn.Module):
def __init__(self, dropout=0.5):
super(Cnn, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
self.conv2_drop = nn.Dropout2d(p=dropout)
self.fc1 = nn.Linear(1600, 100) # 1600 = number channels * width * height
self.fc2 = nn.Linear(100, 10)
self.fc1_drop = nn.Dropout(p=dropout)
def forward(self, x):
x = torch.relu(F.max_pool2d(self.conv1(x), 2))
x = torch.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
# flatten over channel, height and width = 1600
x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
x = torch.relu(self.fc1_drop(self.fc1(x)))
x = torch.softmax(self.fc2(x), dim=-1)
return x
torch.manual_seed(0)
cnn = NeuralNetClassifier(
Cnn,
max_epochs=10,
lr=0.002,
optimizer=torch.optim.Adam,
device=device,
)
cnn.fit(XCnn_train, y_train)
# Specify the path to save the model weights
# After training, save the model weights
import torch
# After training, save the model weights
model_weights_path = 'model_weights.pth'
torch.save(cnn.module_.state_dict(), model_weights_path)
stroke_width = st.sidebar.slider("Stroke width: ", 1, 35, 32)
stroke_color = st.sidebar.color_picker("Stroke color hex: ")
bg_color = st.sidebar.color_picker("Background color hex: ", "#eee")
bg_image = st.sidebar.file_uploader("Background image:", type=["png", "jpg"])
drawing_mode = st.sidebar.selectbox(
"Drawing tool:", ("freedraw", "line", "rect", "circle", "transform", "polygon")
)
realtime_update = st.sidebar.checkbox("Update in realtime", True)
# Create a canvas component
canvas_result = st_canvas(
fill_color="rgba(255, 165, 0, 0.3)", # Fixed fill color with some opacity
stroke_width=stroke_width,
stroke_color=stroke_color,
background_color=bg_color,
background_image=Image.open(bg_image) if bg_image else None,
update_streamlit=realtime_update,
height=300,
drawing_mode=drawing_mode,
display_toolbar=st.sidebar.checkbox("Display toolbar", True),
key="full_app",
)
# Do something interesting with the image data and paths
if canvas_result.image_data is not None:
#st.image(canvas_result.image_data)
image = canvas_result.image_data
image1 = image.copy()
image1 = image1.astype('uint8')
image1 = cv2.cvtColor(image1,cv2.COLOR_BGR2GRAY)
image1 = cv2.resize(image1,(28,28))
st.image(image1)
image1.resize(1,1,28,28)
st.title(np.argmax(cnn.predict(image1)))
if canvas_result.json_data is not None:
st.dataframe(pd.json_normalize(canvas_result.json_data["objects"])) |