From d3a4e78a415692b2f4a9298e788067d9d294bb5c Mon Sep 17 00:00:00 2001
From: RomanNum3ral <116513200+RomanNum3ral@users.noreply.github.com>
Date: Mon, 15 Jul 2024 16:52:27 -0400
Subject: [PATCH] Add files via upload

---
 horse.jpg  | Bin 0 -> 976 bytes
 predict.py |  90 ++++++++++++++++++++++++++++++++
 test.png   | Bin 0 -> 11976 bytes
 test1.py   |  71 +++++++++++++++++++++++++
 train.py   | 148 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 309 insertions(+)
 create mode 100644 horse.jpg
 create mode 100644 predict.py
 create mode 100644 test.png
 create mode 100644 test1.py
 create mode 100644 train.py

diff --git a/horse.jpg b/horse.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..e46f75fe141b6911fd950d057c2018c937fe9dc7
GIT binary patch
literal 976
zcmex=<NpH&0WUXCHwH#V1_nkTWcYuZ!I^=Bjg6g+m4ls~os*M;i${c)hnt&6Qb?Fz
zL{>^(PF6}rMnOeST|r4lSw=>~TvNxu(8R<<Uft5x!pKI?*u)5A2qPyaCl5D|1TU|I
zk)n*E5y{~H0R}-11_cHMW=16jCP7AKLB{__803NOWMu>c1}I=;VrF4wW9Q)H;sz?%
zD!{<R$jr>d!pzFb!U9xX3zTPI5o8roG<0MW4oqZMDikqloVbuf*=gfJ(V&YTRE(2~
znmD<{#3dx9RMpfqG__1j&CD$<t(;w4-P}Dqy@EqR!@?sXqmoln)6z3Cvx-Yf%gQS%
ztD0L{+uA!iyCzMZGIiSY88c@sTD)ZGvgIpQuG+MD%hqk%ckJAC=<t!F$Bv&kdFs;T
zD_5^wzj5={!$*&wJbm{3#miS8KYjl4_1pI!KYxMz#mK-6@fMJTcnr;7f<S*Uv9K_+
zu!H=?$W#u*%z`YeiiT`Lj)Clng~CckjT|CQ6Blkg$f;}`^g%SK=pvVxipfLOk07ss
zeMX$en#l4Q++zrT-D2QjW&}navmk>#!@2&0m+TfM&CB~;aYy-E;g!dq)~!FZt7P#}
zrcW80I(IS)PU3h`w`0wH!HSt%q$}R3Y?yZE#P7NbYz9Z<=6yYrS*Lfs<V0Tj<vgCH
zQ&TIR2mCl!`Rkg=%2R79CmS{IvXwq!wQusn-8)rx?$o;Ey|*s+OogOE-PE&pD$Z_|
zj(om#>XJU8op*NFe9}IbuUooa_?40W^jYu97@jw9zOiL3k4kS?esRe>>&~c0Hy=u9
z72MwO_~o^2D^>k8Jzx1;342z5bmrAP)`>!MPF=TU{1P0#_Mfct-`V+v*{l5d7JXPb
zW1pnNTN%b%vmfho#b>X~E=&n?&MTMbPS}2FZRX~Cwd+spY>F;@VxYreQ#ZBM`WV-l
z7t6g)T}lgY+Y)`+UMeea@8Ou*X(|(@NU&@;@Gtb**%jxxr)C<khi09-@OHi1_BDPd
o4xbh5m@mPzYT2?YYHiPM-mfrHu<j4i(lS?HIX6mk_nrSY0XV9J&Hw-a

literal 0
HcmV?d00001

diff --git a/predict.py b/predict.py
new file mode 100644
index 0000000..b2f04fd
--- /dev/null
+++ b/predict.py
@@ -0,0 +1,90 @@
+# USAGE
+# python predict.py
+
+# import the necessary packages
+from pyimagesearch import config
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import cv2
+import os
+
+def prepare_plot(origImage, origMask, predMask):
+	# initialize our figure
+	figure, ax = plt.subplots(nrows=1, ncols=3, figsize=(10, 10))
+	
+	# plot the original image, its mask, and the predicted mask
+	ax[0].imshow(origImage)
+	ax[1].imshow(origMask)
+	ax[2].imshow(predMask)
+	
+	# set the titles of the subplots
+	ax[0].set_title("Image")
+	ax[1].set_title("Original Mask")
+	ax[2].set_title("Predicted Mask")
+	
+	# set the layout of the figure and display it
+	figure.tight_layout()
+	figure.show()
+	
+
+def make_predictions(model, imagePath):
+	# set model to evaluation mode
+	model.eval()
+	
+	# turn off gradient tracking
+	with torch.no_grad():
+		# load the image from disk, swap its color channels, cast it
+		# to float data type, and scale its pixel values
+		image = cv2.imread(imagePath)
+		image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+		image = image.astype("float32") / 255.0
+		
+		# resize the image and make a copy of it for visualization
+		image = cv2.resize(image, (128, 128))
+		orig = image.copy()
+		
+		# find the filename and generate the path to ground truth
+		# mask
+		filename = imagePath.split(os.path.sep)[-1]
+		groundTruthPath = os.path.join(config.MASK_DATASET_PATH,
+			filename)
+		
+		# load the ground-truth segmentation mask in grayscale mode
+		# and resize it
+		gtMask = cv2.imread(groundTruthPath, 0)
+		gtMask = cv2.resize(gtMask, (config.INPUT_IMAGE_HEIGHT,
+			config.INPUT_IMAGE_HEIGHT))
+		
+        		# make the channel axis to be the leading one, add a batch
+		# dimension, create a PyTorch tensor, and flash it to the
+		# current device
+		image = np.transpose(image, (2, 0, 1))
+		image = np.expand_dims(image, 0)
+		image = torch.from_numpy(image).to(config.DEVICE)
+		
+		# make the prediction, pass the results through the sigmoid
+		# function, and convert the result to a NumPy array
+		predMask = model(image).squeeze()
+		predMask = torch.sigmoid(predMask)
+		predMask = predMask.cpu().numpy()
+		
+		# filter out the weak predictions and convert them to integers
+		predMask = (predMask > config.THRESHOLD) * 255
+		predMask = predMask.astype(np.uint8)
+		
+		# prepare a plot for visualization
+		prepare_plot(orig, gtMask, predMask)
+		
+# load the image paths in our testing file and randomly select 10
+# image paths
+print("[INFO] loading up test image paths...")
+imagePaths = open(config.TEST_PATHS).read().strip().split("\n")
+imagePaths = np.random.choice(imagePaths, size=10)
+# load our model from disk and flash it to the current device
+print("[INFO] load up model...")
+unet = torch.load(config.MODEL_PATH).to(config.DEVICE)
+# iterate over the randomly selected test image paths
+for path in imagePaths:
+	# make predictions and visualize the results
+	make_predictions(unet, path)
\ No newline at end of file
diff --git a/test.png b/test.png
new file mode 100644
index 0000000000000000000000000000000000000000..8bda1118c5a5d9ae8040ebb567c2b9d8e9d109b0
GIT binary patch
literal 11976
zcmX9^1yq#X7Nw-S1tdop0b%HnZX`#M5{C`}Y5D2y6cA9lQ=~)zha6H`KtQC25)p=O
z-gUjTKF76SxZl0!?o<2hSbbe}GGZoTEG#TCO%0VN;D6hH{|NEHr$NBOF)S>BWK9)C
z1D~8%xmwB0n-h_G64@iq52bS^Y}GP8K)fxOV{jqW&Nw#d!72I7cmv#ZhYy`i7__Vt
zqfQ@v&0bgv^L32V#;1ggxXrqDJbV3ns?smJYw70X%YmCDG>%gh4$&kIQ&6OIXq<zR
zD$}VJ%8b$V*MAqCUUAbepM*D=K5Ii5r9!?nXkT6<=64OjlWLGr+6^-(p3wigV8|E+
z{(RT_`LyFRvoHKws_+ujN&!2eL8ZalvP#41^luGQE?*{hcWb=A=C*Ad@65D*{rWVk
z%<;>--_cfzbZ!!OvPR0ldNqbm;k4chMNqJO1OMcT+Z)L4sXut`C>Gc?R{FJld|X=O
z*i^OPQa)m5?X-LB_xso{Zp~OTb!)O*k10_tOSk-|%a?D5W;KRd;CY!<HRt>|6|4+S
zz|!w1%9L6H{#<5Ej44s{y|yT8a{6Wq<N^4;Nyfcf8dcsnB}K{#el0pxOEqGYKAiSt
z#}eO=@UHix%f<zlXBnko?D~0|;5U|xygE<8kJo)&YQYgr1v0&BV5H9~4TdF;{a!Sg
zjwIYwO<9%v;r(^m<;zn_yvM03>5o-yO@5=IDYCS=?Q#&c<y{s9(6}kbCR)`Lro<@T
z^f<81yiKr$Xq2KV96BUyYf=W$G^n123?_V@a|X7ony!L881sCXsYQm5+w1-Ttl<5%
zX2q={nN1#??Lm_Dw4<2kfT+%H*3PkM_K2E0OZ9?F^@PLL@86F#Q!8e6UVfQ>*<|{l
z8#NcONV?a}Zfo*0v942QtZ9Dt6QXuRL<mv(1P-VDj1Bf+YXZ(#Gvjn9hZr;Zx*)Fi
zN|fm{94v#V<s(%-_LBuK`#Rd6%m!@8ofi9fzUJhz?QW@Ib>*~UC8D+(Q9FPAjF-;V
z#18&*+Y3P88Eqs$z_*5fuxXqF2WuDrmj!HhBA2V%Uli>REcy8}xF1oShZjO#8MXi+
zRnv~(e%lawU_-#NDND`#L&N?p8jMj@hShd;k4vrc9u5XEyyG%6(PxZ0ku9QL+wgTk
z4B%5mrg}-TkeU(({&4g4tyC*61W$ZC$->f-lX!L|l=#(`=l947z$d_#X^kCaT)kU^
zWYP$<#&C$`6O>papba5WG1Q<9EE(V1{bmiA@H%LTMNH>uxGPI~;F_FKb-L;kV5|NS
zyQe8jonR}jpO^N4G1}f}bbE%g!w9WByPRE!an<2AXYF&&f5a?mBv@MR4+$Su02oUq
z0VjL)xMOK%eZADNDN8_5@c3{{eg^>Aw06cjw+Q-aX|Kt&2H5G7X$@_E{l#^Amv$-y
zKLwb=H*fx8c7UDm&uyZbvTvSK8$v7f$4(RQ84H+xDuj3F$<X=}-@I()6oNTxi?o22
z`CVtWBIpyyApJl$xCHw6$c@7m%h%v0So$~SAbixGw5oQmN!Jfa^au9VFbc%5K{c%Z
zeDikPBpLDkocaSCsKRabCmcj+qY8iSeqPU8_jsuvWn{&#Uof+CR}}bYyl_GYheju4
z4_MK%e6sTI-rj3TmbZ{t@O4$W78~G7I!~KSA9_0EB<!qG0*D{%cUs+lDg54{-q=Ap
z*AZOH<aQ3@om!5oaCZF+Tl}X!kFv3^Fp=<Y4WEKNMLOBl9t?e-C2Q8!1^xu0O_Qk(
z0|Ue7uASqb9w@uU`hdCxmuYA5>0NKKJ{V-i`o)KQ|HZviH3fFahySE=O<nG@rK+D?
zUK7=`C;M)4pQ3$(M2TaY=bS;5U}CnY=XYpBFgpXx0?*be(%01^h7mFj?fKdD@4%8I
z;_L-L<jWrEb0Ie*cm$Bk151#YaTE)mnRc|*$LHmzC#f;y8wyqtg>aK9izF@KT8L><
z19Y7SeBAc(Q#wzBz_ZFS9sw`D8@YFnJUZ>ybJ|GGI`76ghp!7&z+>i|e{WBggOKKT
zfd1!C;5vO|et-Z{W7tTBpLeB(Yq`j5OvtEZK>3<K4~AcLWyL|nC_d&nbZ-hi_57;K
z(zGT6#1ars0iZwt2hQ1K%>PIFhQj<8u708-So6Cmh~^Rlk0%Wd&&>d~+Yn=%1q&|o
zttL*TPwJaIZ~tg{y?$Nl*XEL7+p+X-ORK9&f(Rr`@P$BTYg?P_!4_0_4<m*34u@+D
zC6Z(v&sY?e%`q3csp87g|M@N>&YB`|jGd+(iw9jDWKf(VjuZIUh@D93OUIP4pYD>2
z>Zx>rkl*ZtAfQP-Plf!;0cyo5AWxJKtSYTA_Sdh&awK)0`U_s>W@j@4U_GQK0T)vJ
zEmZge5>cBDi51UI3}YYnYZI>pxBGL~n@6qi27_GiXV*^r3EWGJiA?K^F>2-v3aOg)
z*<;^$1-EH=1tK57SiNxtlo)&gEI6|1DVD7Rg4ht(VqplKcxDq^Z@Fg;a|Yo&6M3(l
zF#O5%p>;iy0kVxm!~+a}{P^*cnZWV(B##a^4Y=vcNDzv+>LnLLg^>McRw84bBKQ)2
z(jG02ao2(P*+nfcL5O!Cw>i<pRG7c{nj&frK!&sdK<$VBhK4Y{V*~%AITe}UHiVNX
z(@VMh=V2X4<dk3A4{)%%JUUOmio>@`xjCDm#HtV{PU|Mqf5BTk`{maKN9kOq2p?pm
z3uj}UJCYptv+2WL_kgg{U!Nq%Rq6XMs`rm*DkG41&tzvl`HbB@+fQwt*tN)Ee5<23
zIps)QOSH;@<|o!oNFrX|q2i32qKqn3g~sMVQ%nA((T#P3dglv|mQNZF;#MPW8Ytsr
zmcz-w%k3OQ?};aqv6oY}BAmFM$XM=6-!=ZT!!AxUdms#NL(G8K(>V7j@#M|7<5*%s
zGGjeCY;ed%D^X2qK4>ULGHH6-OSd2=zuvpnn6k#5g%70b43;_4s$aX_PHOBam!>FJ
z_S!YE>mc*+u1=Xe2+1q_lTlmdgl}r~g(^j-KGdS-gc;)w2}ysa9h;`WVw``>mJaEf
z8LWgDN8uTzAOueAf;i7V^vRh2?Pg2s!pGg>BmFfM003s5)DpvL+oD%0szk~lt#Egr
zXOA{$gIv7@^Lk^7O}L*Ic07h*+Os*7H9Hp+XZv%XI3VI|>9D73$5-y#8xdsq3gZ>X
zkM(9Wc%i?KTYu~v$G+3_Dtig?6<Jm2YZaq$f)llRVk?Cm6l}_)C4x+i`%Yb`!CBt@
zkV}wN4sZ%Jz8`A}kt(nKSrj3U4ySWPJmy(NGL3NKiXmi<pTtNMEqUq5#u^Y3)@^yd
zKhhr-&(^7qK5Mr(hYhh40`T1*go>WVlB+9{-z{`IPN3r4XLCrri7%y9Ed&mmw<<(8
z7Cy8V>7&?WDuz33C-NNT9ODoXxg3g!<H*uD_yt3>EbWavm;F^wf517uIsKB6iWr}D
zub2$?us3fRBFV9fmqUBzES^mq)|q9KJ+cuI>#@fwaDB|h-%HV!ph)PyKI4iLymqV*
zD=0a$(^Q|(z)#nYi9Luj(DmpGM}`>FMuHHu&LMb;m5!(i^;mt+j&0!Ev6PL`Y*nDS
z^RXPIA^8?p5+=qtp%}FM6T17Blo9u36kamGaE<k%RIOlVAo~`w`=sh2^W?Ud_J5pn
zAoDygv9~>f#*pV1@7sA~Ml(0x1a(N9FFkb-r#V~qlO&F^)@eV%OvG>mLQ>fRiZ(&%
z)0gFrwoi!b3OQSFu~&(-*uOWXTkh$_OPw7o^rIKcr(;h%HK#(+hs`F~g%pL65PpQe
z&m4N=i9mH$RJa<GjIW9bj<w2^DB8%NQ)oCS6rfGGJ0E9uxF4<{Rb<|%%Ho?~g;6<&
z)*Z6^nR!t=lT0<x-4?d<0G8KV_C}<x6#{}5AA2ZcY95<Pr@xGt+Vg|fceNLR*K4-!
z<{=)kd!eizZDG4JR8fUul`>AacD_2mb<A))*cXaRJOyE@IelYyOI|uQ<wJHzu&Wik
z@a9+WLz&ItY1SZYynI}50?NDU>2C;<D+|W679-8;(~3H(heBDmD7Ev8P|hFB_}=K1
z5e8>OALg1PL3`>f2|BuY$T#|=dOE?<uNS3t#2Bw}cu97j`$fr1zX-v6bW4F=j@$$4
z)^?7K;95|^E$wAzg8Cz^j82n+YW&dafaK4BM9VBxQ=F1YH3d`<s*SbnW3dEzLmSJC
zo1DBIJ6d?HiN8q5iB0$AI;m&&E3ndzFY5u|CcZU@3RBl_BScm+WpaEBSMTEfq(-(c
z?d`0tcGHLF;d>XeQrapM3PNRN?Qq?3JQ&W2o^|pe6@->fqmD__2BKcp7fLb((g==S
zp3XUo|8QZ?awA=eJYI$Bj)?6XUjz?orJVALqt?U(BP?@xq+E#MH8VSvAk<DDr@;`^
zr*sczN=D)_=XdehGwUAshz}3U5O01Ht;Ou4UdJeGS1>(g@d&=WUX+f7X&`VvQ429<
zY4n6zykVrG#qdwYB+t~0N_=B@?JRS9Vww)2*?1?3d1#P~k;GzQy&2`cBxaQ<FJE*P
zPi#^Yk_|lWvqMpy>$`-qB%pIi36CRLU)djU*DB5)aj7C}jcTq>uh=7bI6w4W@CM-S
z)E92@ldi1{%56{3g#6@Wz>>(rTE@y7Imi3cmXOb*S)^NTos*x9>w$FcmV{byyxExf
zFOP<UT@HSEg>P-jPfv`BCdz_n@>85$^iU*VnKKfWZ~7GwoZjR&$qf+YO`qLF?xp)!
zkBcn7cK*cX!1XwQ;5#Z#aKJcTLc4sjp3sKOM~)op5O*U6k`UE@`M0oxUUY_;X`rm)
zIE**0Rzl!#^3Of;2rSJTi{yZMU2@KDPVp_XRC`NAlvMtpRUtp#tGK1(sFMpuwY5t*
zDu2^&EOgnwUrQwz&_*^j$X@x4rP~}Lv`zekacxiZS1)0r1u#N4d@sMdv4If13D<6g
z_UgE&PfN1lqaE%72bP59Un#NC*fu!GxN-R<A8KwA^kxv(9Q#s}nLuXLHfUK0X&w2P
zwUaPp{7{}r9@+r3FvPzkyMceC*z!GXr^4pBD2Z4`bihH31bu^?j*#&~ub*%6x`|?>
zT>}qfQ#oQsn4>f}Tl@r1zMq4t(Pno&STNceT4dR=v~zQ_SXB6=^4-U%Z$mrRvS;44
zwW<}vbNYJ|sW*JzedP1cga~U331y8`ui6Tj{I4mYaxCQHznS_mRPD&Efk+v?!&yD(
zf%HCeSua_h|Da=&a1)&%IzAC)4B8h^tgFHu<no<B&D(~UXUZ5V1$h8OOpvJPit-}%
z<=RJx-}>a7GTpm$TYL}cD3m;-mxA||lo^%H`EZi;?E8=STfFsa3u4*F-#fpv<g)H6
zo;HBR<~w<aGyK%7sPi`I_DiO22xvxu>pB~M<kyNPUH<T-(gDWEI4ndPZH-mzoG|kd
z6#ry*8wxkwzLh=cylAoI%=k%Yf^Xla=<WvLDT=&~`oP}x=k(GV<MZ*h&TR=TJ<}(1
zuk;j$XkS>K0be&nUETyohb*jf!|T2L;dhnV84(J|0|_RwfAONp@0&JcxmWuqfnov1
z=fibSty;0AKh|7lT_<@rL<vGBkA8P6W~Z8m{oN#-LBZYPFBf4L!EU?**2e(3eV7-b
zq+aS56gAaue1HO*W-tlr!T(-oS_9g-<WJbPyWaE2d)EmXn{7f7wG=SPx^Ms5o&dQo
z$4_>0CZFAm%u?YWdH#5Vr73Qp(e&4r%bbzxFQ*<Eet5Bm_G;8m%a<;M#9PnH+kJaB
zNb$pIKJNCzp4VhVp$bdZwx|aeCe<ryA<V71sXgRtFvacMXqld=vhtB2tQ>g-Uc8XW
zr_TJk>a_mGz2n?WY7<0ee)qH5EcLD(cp)*rb59+cOvNc#-9`4@^JgXJLyu$0%|k#7
z02<Hxnnk@<81qe?=JTKqS6-1tWxNYTVkQT`semK78B6cBQ0fe=e=lK4z=8NN4f!in
zN23ZAc}IB$EeJhOfluTuow*I%_}e>{_8uy|Q|;9MfL#<C-sR$+71J8ZqV(i&u`e4}
z(|AUiAs$~MGzaw2;5|ORAIUFS30vRD=WU1y@H(`<ogzUC91v+IHu(_F?VV<g*coc1
z63j7r;4KDAvLt+8RZ)3lZ47Nwc_2oiPt;RH-q?@6Puhs6P3}`7?<A3eYWA>8ya@;l
zUI(2aTBxPlDdo{wV8JlSdmo-WR@D?bAD<Ctagi||tCy8e80tAbtjj!U`+$-O6`ZA%
zWloGzE1rlg*orqtve{z?6?_AAQNfS+8pg#TZOK-)u=1O@I^o1y&`_z88c^$Do{+}l
z-Pkq02ZdcdY2UNMtCi-@gM_z2q^zTi)Spr-R?$Lgh0C~UZ9!_Nff*Vz5-x-4XRp65
z&`hKz?s$K#n%Oa}F%<7cHwl=2$w6>sp2whlae$FI<8|zlP1w$<RI}2s+*#Nj2szux
z<>!;Cgk%%Cv^Vx=KdbJ4YOq<wA0MUaOT;|f3q^8|c|J1VCQGJb6kLsxIXrlXQ^7}!
z6%0roH%Wt32oHrWI|nKO7R1tjBHh8tGDq>&KlJ;IEh0Ib1IPChHu6oD?OaOW#FTv>
zF>-xFijR;PlUNHbsn^PKcFP2GGd}|4W?a5BZdodH8citlws3FUr+&vwr8u(osBHX|
zO%Qc;`Kb!k>E`>inHzmvjT;q)R|w8nnV88g$XGOrgPoEthW0_2(}`c(8b785Ib6)Y
zbtEK~%wzGA*R_b`T`_UFZL4X`hb|QlmnPE_8$u22mN;UDgO$4sZsEp+Dyz`;otQ$6
zPLmAKk-s|{b&{mtj5N^vF}%76I!R~oN3ja#!wux}UCI&XVN~U7N3Q{a<9VL`c-==X
zk^_r>gbIH)G*it20pW%61J>;O01j-`f69ko2C#z*E;Ei?q}%g%e&OMg#uPJlF*Kq1
zjvg!~kA&ZRnSZ4<fSv#uEbf}{nMno@Ctq~1jOr=6Uji?X_4pMrCuxL<KxUO;GFKZ_
z9k!vc%m$y0TQ*7kgO^!9MTwX<?mM`ohKtonqe;Zyc!o@=AY?a&YAN1FS73h))sRW!
zn%4gdC@}~Of2H+`%>gIO_n~hxTH8QF8qPb8%^`qj>5GMx$shcs3yonWM@=5Cd`Hm)
zl|-lL$`9g9T41-24g0H3xPM5`=}i%J$q{6aJw3g?x8!g*dEs;a4qsH4XA4LpFQ3Uy
zY(2AL0rcA4`5q|527#I)PtuO?!$EH{BM~v5qp7|we3T-kx%j%EOXZC1?ICwJ3|{wX
zuQG=RA+Y@DEt<5-dlvtmGG_?cF;|68uHdgLhI9$4^=%oaj1reGPF}ZH!yS7s0T%<>
z&(~jHaIsR7`@z`5yrlOw=Go+gLnSCK87Y*)yl`_=gOGimeo>W{1j~>-)sp4d_Cx&k
zROS$QB%@@=7)>l%2)T614Q;*B4|CPDRtcV^;g8m=r|MKmh0vD*HVIHAwNVixKP&xX
z@Hm+&Y#*wK>+TlEpecxo1_Xw_PlW>hmC#$!>TQx{swWFSQ*w{7*)JZM<vt!S@o6@r
zGxI!hXA2d`EBQ7hj~*^dk<Pd=EfvNT^|#*%FzT{r3#b(f0=M^6`&`Y5YN)?MBpEky
zvRsdh=417LQlo`F+N18@plP!E+^s(`Q|Lg9WeAFW5>D6hmrK}a%h}OwbMH)Rybna>
z5{A^=o?cZN@MbW(v51Pz-fJOOK5v9{#tPb5Dh)v86qPw8nO4-(Ai1=y?Y>+7+UP34
z8Xx(6nV(C8b;D%1{RFM+Vo9Y!3JuD9B=QC}xaPH_E0Zux2b2=?j!hr@<jZre$YbAk
z%$c<}!-FW?>HIWlg?)+mJ+=l;!{3-z+y%7fv)Q1S9w$qMW8aUcX2d4$=mX{hnHk_N
z50EC-=>Ee(|Mp2W6#hF4_!OlNwO#?1@aPz;l>-{Tgsgk%5uQYBiV!kd!)2Cj+PXLt
z=CXLIFnU;AUgyp-?+M*GZUrRTt{D+IR(KIHLDtnVlDMmhU0mi|FSV(u%`?)n!lA{?
zQHGWI=#4V+@A`aRkWi;0C0<z|#*@B}d)eKYbFGk{*iYL)p_=F=dx}P60T#TIHwo*n
zKWbm(hvKJmoy0Zc1s5Z@SI$;6$K(D(2=rci+V(s?J24@ie=(ny3umRZB*kgU7#m}c
zyjZBtL3;C&8*(^Wz)dx*M}ORf`eXh&=U^dIqYvj8j}DTlK~VrF4scJZ+!6_Ps*sEN
zXQYG=qv!zP@;u)a+eeciCj8BrxNt2;t)e#g+Utkz(Ayv?38nNnk8Qo_RIk8_u;3g2
zpcEg9rRc+attE)JBh&1l@A}?*CMilLa#aG;33P?Yl~Xg@^5_oSr`gp{2;|9(w&nI&
z?d5yBWQk+S0uZhtVX|oDM=ia@6g8n+W53~DtA#l*t455W1qsEKx!0TcebsDQ=%f7O
z-W^{yw2uo*bm90sAZKXr)gFFNkf<mxenV+7J4PH%jLsh#($Z|;kqqS>7So_elUmeA
z+b7*CdhorrYwLzpa*An%NMEGlfhpaO(*|+LJ=smO7*u$tnT8Ua_J<|PA4VKpL7~yD
zQW8y|l!xv5o5})363p2DC9zFJe^`?uj7w?T0#?^LNb#1^_m^Oz6x)O(>;C3uyM>ip
zeqoZzqLV}XDqJO!LJJX1B{VAatS=tvv+0SyxS1vxEYhz!p?Obbf}j!e%IIVG<K4=p
zy{wA?3myGH&X804dmkdnAj0eqrWj)~RY;Tm6cgNYsW;9UuJYadWUq12Z%TRJ!IG4C
zf5Yd|nsoa^t@L>8z*ovkQKu&!;%XQiE)Ly;jd!fnq%#qkmMrX<(7yT@PVuDMMwfh4
zjLQ;%C0nz&!DoIg|BphK-Pet-bq1Vo33gx<GuvLQ)UpJY0Tl0p2s^Mo`=5tbWtc6`
zk>;(U-Xm34`uZ*qF}3%v<`nU-X7P}tr^~T~jubi3I?Jx&zp8ONYchYd$jy`mtu0<D
zMC?&ms!$HhxG)-|4oX=vnwF>b8KDWviQHY{_Ac?jkf-K$Ia?<}vN=IW1`jhVXcO4}
zA@y1h?ej9qj^>3vR4ngL8sDjc)oBYVzD5$8nhfs$*G%7`w#|2iPFyuf0bao3Aj(9M
z93RKZ+O9m$;Be?95T}&VA;jjrsT~)yKlqjz!<0<nBtE;NS4X*=Kxm^pXwM?`ZVb$T
z)?x)Ty#*qrBVOkkT_Nl+f5;`3HTS8)JyZVE=rN0SC}`784sM*Ajb=Z~fE|(qy(KDd
zZf@3ZLwwoXD0aDxPLfY>XsEx+4nU1y8%+tfQ<tu8CT0$mE<pCiA_#|jCM^=iUUN_f
zeHhRSeKQ-%QY)%&?>iShvJoirNpXAPPr%rcx_m|F<r5YW6xQ(5CD_>)M6Dh^%JSXO
zD@h|Fc+7bIK7H*YMYf5IW^CAW2e~WGO-OH-g++EQB$g?$Z?HJyED*ZULOoph5dHc9
zKS{a7*O7$9_eEk;*iKGpNppl^DosH}C*%$^C7h%1SoEryykXKyTI?bo88(U4%NX^E
zimC`ky;q;g&t^gCNG(2Qgo20&bP}f8v&%`XVhzv+Z6pl*@sjqtII1jyeg{>dI`%CG
zA!1l1u>i9kuPoTR)@H_rqP*80x<0?EB6&Js1BohZw&iAm6!aA?7S*c;4K=JyKIEj|
zx?ESFr!Q^C+8mVfxif)x$^PG``B!NM`=;!%k*P*9v^YlTk9*}7)KlSBw|*?x&1?QD
z)PG_H^JxYjF=FlYYla$AOZM{TH}vcjHU_M``$H1vUt=iJAJ!^DxC-Z4W!t=VK@L<`
z%CIsxZ;phq!%U+K`W=EnB*g|i4)ZY{FHBCe`AMa5HvPtWwP(RM8rt9ddXaFhiqu$r
zM;_!oWvUmFB<Zr0n{z8tWT*HixDlN)JcNLX*SUT3B>?^Yv9^7(N;DPHKxuT}8YYLC
zFeivlgVY;V2Yb=H(!U*)btcwVE;@N5nmxiw|L1}mVkn;0oi=hf!sW5tBGl%LrgKe*
zAKuGNi(*kbUVwk>5A>}+MvjEwvUkk}!A4?WM7WmNb-dY{ak^}ec4Ud_$$fE}4`QjW
zPU26Fel|)bq}slJ4Q*4@KsrN56%M%Nv1?Rf_x_`-%nNN`MGV>DSbs<%>%W9c4tN%j
zpLShhEwy`NCRo7X)7LRl-H2dJ^W#N!(ZC7R&bEaIs(Si2<Z(covg=ZO;-VX%q4Ey3
z$dyMGe;179Zk`WfaoF1i<gOH_yta1sU>#k#qY|W+JjljSxTTV|b}ID2zE%~Uqn)xU
z`iG~sm@zSiI!?GpMkcl-*=F=M(28$pO`R3YIJHiSCo}-LgAU<cTvj&>B%m<tyZYgq
z{g#X9d-!;4ctP^Sa(8Sq#p0UT*F$0Ij-j^mA0nUTzOcxY@>nX<`qb18fqBP0p?<<3
zrD@Z7FfT%_=%ZG~9dgP$x8R+38;B-UK3d@5rN6g|%ll-M66oLQT4(eM!yiPTwkH2=
zkfEHB7KQPYzfzzZExX`+e|m5rr^eMV>%f_*qxvGbXdox*hEN{2Y*tTCU_kv%IT){Q
zgYnJdRI~~NMd>3Vu_=_PPw*I$yDUOFYBT@<ir+dmc2}twZexOUr8^q{bJD)!t|LHD
zHb}0P*lPy)rCuH#u@FAu!DGamBOEvy*Lc0x(GB~g7=O5fRa!!P0%IbN2IC=oW{Z0(
zLxNpyR}OrO+sW_(0t}VLvH)`@S-lUhOF^Uv|1X(lT^Gy&xrbC_dFC9(Oa2~e;H{Cb
zA{GKA?XFf3Bj)?~&#}1tFk#oe^Dsp^uKOR7^Hl2xRAb$_y5&5(76`0z9w`ya?rc=p
zmi1vHprCvb?cZo2Q`(H9XP#9*3l|j6%sz=$f4`+3uk@QkX;3Fsx8R355q@_;fcstn
zw%DB{^nC{I$LDu7RyQ*=f1CB^yB9hQK810<{kHApSlW^QpDn5aTf>1rs_YqPo^n;1
zrQh`15hz;hE~8e!AogZAdRT}h_E$2`N@bqc{(wqqOD!|f{^^oWWt8+Eq~<!a3!0d`
zT`9^<0AG64vpyQGX{bNh#2?pP@frvZJ{;MfB}KY(p$W#OeaXgAgjIX}ll&;pBB*qi
zZui{e3n+9i-qMX;5aTm=XN#E>A6PYOY;$9Q_4XWRsqko1iGKp`$mal6Xb?l8X3st$
z&k(YN=|3GPN5gAjc#8XOa(8mIQi)!+V;y9F!%YoTj@Dnda%Mto@Oc#su+fF^zi|T0
z&fr7H5R7!UuQZx~xHDk&>};Gxj}u0JOrjs3aR=>|FGq0pOK}lPQU9I-<!?HRWd9!M
zjYjOz{xFW<JPN`E!Rr3TqZ5~N9Te<wGZEYnmMQ**xl?Z`o}(ugk6Jj3AvQ;Y)86fH
zlMI&dnB20cgkKBv24L1EQZZ&J$=-r*ydoQ-o<qqQhO6IniFrEuuFgE~rF*a$#Q@;=
z@vp&L`@0fLf4@jFKfN}ZY*teI?YAi&dzWtiluU{kO5|JVJ1D;<x{dYp7widTcXe0d
zO#}yP4fz7E&fxRAKkP{Li2{23r4}A>!*sm#MG6haJ6GFeA?`_e5r21hin6m}ctkLs
zV^f@N`8qu4094!^VY|m?7#>2LLze^WL0Y>We-86~jyozqd13&lNT@9;CaKrT{ouiT
zM612F0(Sm?dW}wC@CCFDggD2Sfz@?`n$EnZ-%Q;NDXJ-PYsu@_@=QEjr=()C>!fmG
z4iyn(F%`(dI99FeiU*?UedOJ2M7g+mD$(-U8YMm)Z%)g;{XO=}Xx!>Q?80}ZRHyTi
z#?uj^NDiptW9}~l{CFDgPkF6`nL}hquT>$Ka67;Cz#<qp{8jb3=zYv-qXPi9%9>PY
z;og{qCR#qeXNk1qmxJh^PPCVFlS9b!$%_SN@spRYAm(ooMUfqaK^Cz5S>+D=&uh#|
zdDHG(tn{$Qj8UJ-i?A<<eks^Ej9n-R4<(JU<&D_!jBaj*_|}n8jyk!SvT1JaVg>ar
z6C`a|3MjRzeYRhl#bS+1u=n4^xAg&|!sc$9@su;xXN|Yc{Pe?9R}P{L@2&j~W!^bU
z=VpnkpgZ#d_^|A|*EGjP7Y2Ap1<tnB<$tf~g~BaxCZW^%)^ba%9zEQ}9p_RwOOE)3
z1M`xqDv(YOeLc&&H;@l|A4<0|+7~v$_2mqhKs*g77d|RGR3f4qaC!fin&f_R>bSXZ
zk`2Bw`|rR1E^aROc7JwU9mG#l>uXi{<k6sLiF`v$MA*Z{53TC!)YI#{NjC~lkn~b4
z<U9r$4qS05b|HdGRL3G7+K;U|C6kiXIOg`1nWQXaKF#v<9Q(bJrO~_A$kSjM*LcVB
zf$#$l+pon%AV+yEeb3s~Jm|di=B!`-_wM|c&+hTPO8Ihu-@yfpQ-mobv&uecoCDLh
zDH~xaxa;*6C2htCbD|nka(~n%&?o^3N6~-4fI%419Ay0DipDR7q><|hqIIs2Sf!2P
z&;iFLFbrXnYj0G$4r!8&T)U+V_w0Hlp1tiRX<T%YPrJgy61YCr4fIMMYip0ugETK*
zym+AV6f_(Q3kxO-;gE0|IFBXefO5xD=FvW_EH%uOO^ggqLaKxqFOG4Ec;Y94fXlOx
zp9sZCa^XCL{dJ=6u<v}{I42|58&k3_5r5Hb)HlVc@{Ex+`dbKJH=8xu*M-8KnLL_z
z==&&*%a`Zig<!JbDExl*R4Qn{m%$vWnVJ^;&0>5^?sE7?I1?;;WVjjVBK{55{o3+W
z5hfXG5J33$t<}rsML;d1-s#2K9x`HKJaLZw7C2cU*kToh6}?yG(KIc;NcOf5o$=^!
zvqG4D6Uks*M~U<8>sK25F?klyHFzGDbPy`Be_Gpq0kkmY@>l!e_j?ShOE<^zcY5O)
zfOKhZspHFjl+4v{qmJ)u{mEji$z(H?$n23a7Z(?2L((%)Z39uy2|SvT{wGqQEWcsp
zZpPH3_SSMT_Q7{kD{ioBGTR=#Tgo;xuLgZfTnQVt#JDOWjh-;4{62}ZbWYaDL9uWh
zdzsDbEJRb3kraryruI&c%jK^%L1uRe_`1NGWWDg~SH*(M^BW7?mBa8Zh5X=~ma@r8
z7gyKa;<Fg5Tc1SW*m~r9r<Xt0fm#1)EnZVtg1OjZIPErOt^c=|G<~IlPKGU*PV?9+
z2pQJSOd6R7mICu?H_mCR!nXkKy!RMM(^)nn@mlb$ZN%=t6-RKems7S7jxArmij1K}
zW>;3WBxzsBUmv9bh7c%xjvevCLxZk&%z+~7d6R7Z7*Iojp8zcm*x2#L2s2Q|TrL2<
zeA=<pek>16-}nfV)GOHJgFZCyAKzU2zY=E-yf_*pZUE#9P>%uC7c}nZU|s9_=CosG
zdODp7SKO}+u`}0be}Q7VDS(k4pLX27a9u>EK)#NFuYy-(Q<Ys`0l;ht2A{_R_yQH&
zr#a`N6%+v|x2k&?8xImC3z9^sh5VEm{djlEAD)$7-<*YF0T<_aCeqS#7|8xl1_fj{
zMn*>7Yorg}Wai2G+bK`Dvw*<~!gZ;A`s;st+DYmHw`Z8*fBvyG!NvUQ`eLVcS1RH5
zYEwSZFU^b!h$9LS?E#hYx94ME5bt;VJ3Z*;+&*#zvr+=q?zrA{J>LDY<k({N+x+$-
zb3vk!)ROV{_2nsd5g4|+e5wA2Kl8xj^yvWe08Hc!COH`3Y8&9~hYLY#!6;hL!$-w7
zA&jIr>_v*Y#sRq5HH<OPGJ(r&ZE30J`4w6ND+ZMOy6?wwz5Jgd+vaV4ws#hGW@>jQ
zn4q<AuK=0VXbQF?AjbmIjJwL=y~CTHKsc7YsoC(ec*5Zy=_1OMn{Y=lGGE5GXiWar
z4EWoLwza&J@`T13MYJs!L3k**cMI6KgQ)QI<-x%dP<w)XOl})|uMd6g=~?{s>m!h;
zheX_2p5H^ui$|CWg=xaAeK@EdYeN)nfr96s`G6IEUtE}i`2dt>^Pq4702JJBROuRK
z*5zO%TY{U#`|AtOusP?l`T6IG{npIyj0G~)l=;%v?gkJGPFZ?nXR2syfUpN1K=as|
zm_{IjZxKK*wdGM`IB8Ws{@+rd{d%XHr05nCDLznD&CJYv{tVId0+KoYp`+{5r5*4Z
z@0v_Ob0D4Bie?KMw-FYCSOzn`QmO*-5J@d<RKFGkCRCeravAu~@u0-ub88+5O@Z6k
zHBQWbaTUuRNx<RiHGlD9dwnRg%7EAJ^l$FAEV$FZe@DrO%*@PYW0x;4{Y41t23DSl
zzq{H@E=@7dlkxR9NYR(6rCzrZe!n$cZS)T)x%?+yucdQ;0~P6yy8$PE-h(PjKtSNb
zbU<KW&9)a{(+k1WV__8J2-Sb_3H<J8L}F<RgeNcz0OfOQN5_|~{ap-TZ0&Cw`5%$5
zov%{q6=(5q042~1bMu{S*LHaM<V3N0^3l=uQveYPHaRe|%XY>N<d0JV*Vu7B;(GLV
z;gL@HKQ)`I|A{u3>;Zs*ZpGFFh}ca2qpaHL_tdCj(J6L_Jn7C$g-i@(3IcxzbRpy}
zU36=5zxDOmTC(5eNlo{@PZjaS^BO~7q<@P7@mX8I`4iw@JR^dd;7(crXF#h8K4D@j
zCntu$6@d;H=&{B#&4Vic1>XNUE^|S2|NHmvQSj|)Fc=8;AxokKUS*Z37b}cVPM-O2
z$~%F}I6j^+->D70t_HKlmqw?jr+v*)sFhfaS;wY-2bJ|bQ27_FApd$8Kg?<*WP9Q3
ze<x;)^3BmI%EmNJEcZg%aV1jurGfYQZVGl>743Rg>y{UTSk%A|RAyfP{Z|f92wFxb
zDXwgrIZWjWfN7j23ted#6{p-eou#LE_FN-(sM+!*a72(SmI6<6!2}f15%kJKE)K`f
zI3iLbg5{;cwe3`Nb{lZiRd+Ba1==Ve+Ik|DcfWAUN>R=Tz-8<FaI>|w^*it_s&OY3
z=ZO3frA3S*tc@m47oGTjNi5hIh!p1mF2nIiiTB!n+PCN;Egd2`ek<f{f=VAeWm_IP
zaq3ymy`WU;m3uH+hc)}T3N#Qdj{JJYo4z(_XeoAs-u03B{MQ1gU+p}DgKHaP<SKUi
z*@CC8uiNO%2Cjg3j@lk*Z5Ru`ml4{ewME_8Ikd_z|8>;tPK=MHqzkPAe2Zpww=+10
z@cWoT?<(gAHCNiWuNWU^UKoXL{Msy@KIAPJcfoP41nEiIf^QcYZ7K0G)UW$dmvF$E
zu*#*yJ6^d-0)6B^Z|V60WY9q!wZQbjx0g%wsvzbz#d^xUjU2PJDPtT^;On-2Kr3sj
z`jGt#+<rfBsXijlnwAwO@f!(x9QVC$IjDocV_4$;DnSvle<gB2jR%Tv2S6de4*Zu0
zSz@6o8b8@V9yiJ=^SixigXF>%KuQ8q<LYWYP=YfzUIc&7S@5UIz^wrabu3L)U6qeY
H7GeJbx9ERk

literal 0
HcmV?d00001

diff --git a/test1.py b/test1.py
new file mode 100644
index 0000000..4475c2a
--- /dev/null
+++ b/test1.py
@@ -0,0 +1,71 @@
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from torchvision import transforms
+import cv2
+from PIL import Image
+from pyimagesearch import config
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+# Define the U-Net model architecture (similar to what's in the article)
+# Ensure to load the model and its weights as per your training script
+
+# Function to preprocess the input image
+def preprocess_image(image_path):
+    image = cv2.imread(image_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    # Define your preprocessing steps (resize, normalize, etc.)
+    transform = transforms.Compose([transforms.ToPILImage(),
+ 	    transforms.Resize((config.INPUT_IMAGE_HEIGHT,
+		config.INPUT_IMAGE_WIDTH)),
+	    transforms.ToTensor()])
+    image = transform(image)
+    image = image.unsqueeze(0)  # Add batch dimension
+    return image
+
+# Function to predict the mask using the trained U-Net model
+def predict_mask(model, image):
+    model.eval()
+    with torch.no_grad():
+        output = model(image)
+    predicted_mask = torch.argmax(output, dim=1).squeeze(0)
+    print(predicted_mask)
+    return predicted_mask.numpy()
+
+# Function to display the original image and its predicted mask
+def display_image_with_mask(image_path, mask):
+    image = Image.open(image_path)
+    plt.figure(figsize=(10, 5))
+    plt.subplot(1, 2, 1)
+    plt.imshow(image)
+    plt.title('Original Image')
+    plt.axis('off')
+
+    plt.subplot(1, 2, 2)
+    print(mask.shape)
+    plt.imshow(mask)
+    plt.title('Predicted Mask')
+    plt.axis('off')
+
+    plt.show()
+
+# Main function to run the script
+def main():
+    # Load your trained model
+    model = torch.load(config.MODEL_PATH).to(config.DEVICE)
+
+    # Accept user input for image path
+    image_path = input("Enter path to the image: ")
+
+    # Preprocess the input image
+    image = preprocess_image(image_path)
+    # predicted_mask = Image.open("c:\\Users\\anon\\Desktop\\Summer Research-COPY\\unet\\dataset\\train\\masks\\0aab0afa9c.png").convert("RGB")
+    # Predict the mask
+    predicted_mask = predict_mask(model, image)
+
+    # Display the image with its predicted mask
+    display_image_with_mask(image_path, predicted_mask)
+
+main()
\ No newline at end of file
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..0661d9e
--- /dev/null
+++ b/train.py
@@ -0,0 +1,148 @@
+# USAGE
+# python train.py
+
+# import the necessary packages
+from pyimagesearch.dataset import SegmentationDataset
+from pyimagesearch.model import UNet
+from pyimagesearch import config
+from torch.nn import BCEWithLogitsLoss
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from sklearn.model_selection import train_test_split
+from torchvision import transforms
+from imutils import paths
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import torch
+import time
+import os
+
+# load the image and mask filepaths in a sorted manner
+imagePaths = sorted(list(paths.list_images(config.IMAGE_DATASET_PATH)))
+maskPaths = sorted(list(paths.list_images(config.MASK_DATASET_PATH)))
+
+# partition the data into training and testing splits using 85% of
+# the data for training and the remaining 15% for testing
+split = train_test_split(imagePaths, maskPaths, test_size=config.TEST_SPLIT, random_state=42)
+
+# unpack the data split
+(trainImages, testImages) = split[:2]
+(trainMasks, testMasks) = split[2:]
+
+# write the testing image paths to disk so that we can use then
+# when evaluating/testing our model
+print("[INFO] saving testing image paths...")
+f = open(config.TEST_PATHS, "w")
+f.write("\n".join(testImages))
+f.close()
+
+# define transformations
+transforms = transforms.Compose([transforms.ToPILImage(),
+ 	transforms.Resize((config.INPUT_IMAGE_HEIGHT,
+		config.INPUT_IMAGE_WIDTH)),
+	transforms.ToTensor()])
+
+# create the train and test datasets
+trainDS = SegmentationDataset(imagePaths=trainImages, maskPaths=trainMasks,
+	transforms=transforms)
+testDS = SegmentationDataset(imagePaths=testImages, maskPaths=testMasks,
+    transforms=transforms)
+print(f"[INFO] found {len(trainDS)} examples in the training set...")
+print(f"[INFO] found {len(testDS)} examples in the test set...")
+
+# create the training and test data loaders
+trainLoader = DataLoader(trainDS, shuffle=True,
+	batch_size=config.BATCH_SIZE, pin_memory=config.PIN_MEMORY,
+	num_workers=0)
+testLoader = DataLoader(testDS, shuffle=False,
+	batch_size=config.BATCH_SIZE, pin_memory=config.PIN_MEMORY,
+	num_workers=0)
+
+# initialize our UNet model
+unet = UNet().to(config.DEVICE)
+
+# initialize loss function and optimizer
+lossFunc = BCEWithLogitsLoss()
+opt = Adam(unet.parameters(), lr=config.INIT_LR)
+
+# calculate steps per epoch for training and test set
+trainSteps = len(trainDS) // config.BATCH_SIZE
+testSteps = len(testDS) // config.BATCH_SIZE
+
+# initialize a dictionary to store training history
+H = {"train_loss": [], "test_loss": []}
+
+# loop over epochs
+print("[INFO] training the network...")
+startTime = time.time()
+for e in tqdm(range(config.NUM_EPOCHS)):
+	# set the model in training mode
+	unet.train()
+	
+	# initialize the total training and validation loss
+	totalTrainLoss = 0
+	totalTestLoss = 0
+	
+	# loop over the training set
+	for (i, (x, y)) in enumerate(trainLoader):
+		# send the input to the device
+		(x, y) = (x.to(config.DEVICE), y.to(config.DEVICE))
+		
+		# perform a forward pass and calculate the training loss
+		pred = unet(x)
+		loss = lossFunc(pred, y)
+		
+		# first, zero out any previously accumulated gradients, then
+		# perform backpropagation, and then update model parameters
+		opt.zero_grad()
+		loss.backward()
+		opt.step()
+		
+		# add the loss to the total training loss so far
+		totalTrainLoss += loss
+		
+	# switch off autograd
+	with torch.no_grad():
+		# set the model in evaluation mode
+		unet.eval()
+		
+		# loop over the validation set
+		for (x, y) in testLoader:
+			# send the input to the device
+			(x, y) = (x.to(config.DEVICE), y.to(config.DEVICE))
+			
+			# make the predictions and calculate the validation loss
+			pred = unet(x)
+			totalTestLoss += lossFunc(pred, y)
+			
+	# calculate the average training and validation loss
+	avgTrainLoss = totalTrainLoss / trainSteps
+	avgTestLoss = totalTestLoss / testSteps
+	
+	# update our training history
+	H["train_loss"].append(avgTrainLoss.cpu().detach().numpy())
+	H["test_loss"].append(avgTestLoss.cpu().detach().numpy())
+	
+	# print the model training and validation information
+	print("[INFO] EPOCH: {}/{}".format(e + 1, config.NUM_EPOCHS))
+	print("Train loss: {:.6f}, Test loss: {:.4f}".format(
+		avgTrainLoss, avgTestLoss))
+	
+# display the total time needed to perform the training
+endTime = time.time()
+print("[INFO] total time taken to train the model: {:.2f}s".format(
+	endTime - startTime))
+
+# plot the training loss
+plt.style.use("ggplot")
+plt.figure()
+plt.plot(H["train_loss"], label="train_loss")
+plt.plot(H["test_loss"], label="test_loss")
+plt.title("Training Loss on Dataset")
+plt.xlabel("Epoch #")
+plt.ylabel("Loss")
+plt.legend(loc="lower left")
+plt.savefig(config.PLOT_PATH)
+
+# serialize the model to disk
+torch.save(unet, config.MODEL_PATH)
\ No newline at end of file